Update app.py

app.py

@@ -4,8 +4,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 from PIL import Image
 import torchvision.datasets as datasets
+import os
 
 def load_model(model_id):
+    # Create offload directory
+    os.makedirs("offload", exist_ok=True)
+    
     # First load the base model
     base_model_id = "microsoft/Phi-3-mini-4k-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
@@ -14,15 +18,29 @@ def load_model(model_id):
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     
+    # Load base model with 8-bit quantization and offloading
     base_model = AutoModelForCausalLM.from_pretrained(
         base_model_id,
-        torch_dtype=torch.float16,  # Use float16 like assignment22
-        device_map="auto",
+        load_in_8bit=True,  # Use 8-bit quantization
+        torch_dtype=torch.float16,
+        device_map={
+            "model.embed_tokens": 0,
+            "model.layers": "auto",
+            "model.norm": "cpu",
+            "lm_head": 0
+        },
+        offload_folder="offload",
         trust_remote_code=True
     )
     
-    # Load the LoRA adapter
-    model = PeftModel.from_pretrained(base_model, model_id)
+    # Load the LoRA adapter with same device mapping
+    model = PeftModel.from_pretrained(
+        base_model, 
+        model_id,
+        offload_folder="offload",
+        device_map="auto"
+    )
+    
     return model, tokenizer
 
 def generate_description(image, model, tokenizer, max_length=100, temperature=0.7, top_p=0.9):