Update app.py

app.py

@@ -11,13 +11,18 @@ tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 tokenizer.pad_token = tokenizer.eos_token
 
 # Base model (CPU)
-model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    torch_dtype=torch.float32,
-    device_map={"": "cpu"},
-    low_cpu_mem_usage=True
-)
+base_model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    device_map='auto',
+    offload_folder="offload/" )
+
+base_model.config.use_cache = False
 
+model = PeftModel.from_pretrained(
+    base_model,
+    output_dir,
+    offload_folder = "offload/"
+    )
 # Load LoRA
 model = PeftModel.from_pretrained(model, LORA_PATH)
 model.eval()