Update app.py

app.py

@@ -3,22 +3,23 @@ import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from peft import PeftModel
 
-# ① Base model（基础模型）
+# Base model
 base_model_name = "facebook/nllb-200-distilled-600M"
-
-# ② LoRA adapter（你的模型）
+# LoRA adapter
 adapter_model_name = "entropy25/mt_en_no_oil"
 
-# Load tokenizer and base model
+# Load with optimizations to reduce storage
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_name)
+base_model = AutoModelForSeq2SeqLM.from_pretrained(
+    base_model_name,
+    torch_dtype=torch.float16,  # Use half precision
+    low_cpu_mem_usage=True,
+    device_map="auto"
+)
 
 # Apply the LoRA adapter
 model = PeftModel.from_pretrained(base_model, adapter_model_name)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
 def translate(text, source_lang, target_lang):
     if not text.strip():
         return "Please enter text to translate."
@@ -33,7 +34,11 @@ def translate(text, source_lang, target_lang):
         return_tensors="pt",
         truncation=True,
         max_length=512
-    ).to(device)
+    )
+    
+    # Move inputs to the same device as model
+    if hasattr(model, 'device'):
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
     
     outputs = model.generate(
         **inputs,
@@ -56,4 +61,4 @@ gr.Interface(
     outputs=gr.Textbox(label="Translation", lines=6),
     title="LoRA-Enhanced English↔Norwegian Translator",
     description="Fine-tuned NLLB-200 model with LoRA adapter: entropy25/mt_en_no_oil"
-).launch(share=True)
+).launch()