Update space

app.py

@@ -67,6 +67,7 @@ import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
+from safetensors.torch import load_file, save_file
 
 # Define model names
 # MODEL_1_PATH = "./adapter_model.safetensors"  # Local path inside Space
@@ -85,7 +86,7 @@ def trim_adapter_weights(model_path):
     # if not os.path.exists(model_path):
     #     raise FileNotFoundError(f"Adapter file not found: {model_path}")
     
-    checkpoint = torch.load(model_path, map_location="cpu",weights_only=False)
+    checkpoint = load_file(model_path)
 
     key_to_trim = "lm_head.lora_B.default.weight"
     
@@ -99,7 +100,7 @@ def trim_adapter_weights(model_path):
 
         # Save the modified adapter
         trimmed_adapter_path = os.path.join(model_path, "adapter_model_trimmed.safetensors")
-        torch.save(checkpoint, trimmed_adapter_path)
+        save_file(checkpoint, trimmed_adapter_path)
         return trimmed_adapter_path
     
     return model_path