Update app.py

app.py

@@ -7,9 +7,6 @@ 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)
@@ -18,27 +15,20 @@ 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
+    # Load base model for CPU
     base_model = AutoModelForCausalLM.from_pretrained(
         base_model_id,
-        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
+        torch_dtype=torch.float32,  # Use float32 for CPU
+        device_map="cpu",  # Force CPU
+        trust_remote_code=True,
+        low_cpu_mem_usage=True  # Enable memory optimization
     )
     
-    # Load the LoRA adapter with same device mapping
+    # Load the LoRA adapter
     model = PeftModel.from_pretrained(
         base_model, 
         model_id,
-        offload_folder="offload",
-        device_map="auto"
+        device_map="cpu"  # Force CPU
     )
     
     return model, tokenizer