Update app.py

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
 
 class ModelInput(BaseModel):
     prompt: str
@@ -28,12 +29,15 @@ try:
     print("Loading tokenizer...")
     tokenizer = AutoTokenizer.from_pretrained(base_model_path)
     
-    # Download and load adapter weights
-    print("Loading adapter weights...")
+    # Download adapter weights
+    print("Downloading adapter weights...")
     adapter_path_local = snapshot_download(adapter_path)
     
-    # Load the adapter weights
-    state_dict = torch.load(f"{adapter_path_local}/adapter_model.safetensors")
+    # Load the safetensors file
+    print("Loading adapter weights...")
+    state_dict = load_file(f"{adapter_path_local}/adapter_model.safetensors")
+    
+    # Load state dict into model
     model.load_state_dict(state_dict, strict=False)
     
     print("Model and adapter loaded successfully!")