Update app.py

app.py

@@ -57,12 +57,23 @@ files_to_download = [
     "tokenizer_config.json"
 ]
 
-
+'''
 # Directory to store downloaded files
 model_dir = f"./{model_id}"
 os.makedirs(model_dir, exist_ok=True)
+'''
 
-snapshot_download(repo_id=model_id, ignore_patterns="*.bin", token=auth_token)
+
+# Use /data for persistent storage
+model_dir = f"/data/{model_id}"
+os.makedirs(model_dir, exist_ok=True)
+
+
+# snapshot_download(repo_id=model_id, ignore_patterns="*.bin", token=auth_token)
+
+# Download model to persistent storage (if not already there)
+if not os.path.exists(model_dir) or not os.listdir(model_dir):
+    snapshot_download(repo_id=model_id, local_dir=model_dir, ignore_patterns="*.bin", token=auth_token)
 
 
 '''
@@ -87,6 +98,7 @@ with fp8_autocast():  # Enables FP8 computations
 model = transformers.AutoModelForCausalLM.from_pretrained(model_dir, quantization_config=quantization_config)
 tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir)
 
+
 '''
 model.to(dtype=torch.float16)  # Load as FP16 first
 model = model.half()  # Convert to FP8-like (closest possible)