Update app.py

app.py

@@ -4,18 +4,20 @@ from peft import PeftModel
 import torch
 
 # Directory where your fine-tuned Phi-2 model and associated files are stored.
-# This directory should include files like:
-#   - adapter_config.json, adapter_model.safetensors,
-#   - tokenizer_config.json, tokenizer.json, merges.txt,
-#   - special_tokens_map.json, vocab.json, added_tokens.json, etc.
-model_dir = "./phi2-qlora-finetuned"
+model_dir = "./phi2-finetune"
+
+# Directory to store offloaded model parts (for large models).
+offload_dir = "./offload"
 
-# phi2-qlora-finetuned
 # Load the tokenizer.
 tokenizer = AutoTokenizer.from_pretrained(model_dir)
 
-# Load the base model. (Assumes the base model files are in model_dir.)
-base_model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto")
+# Load the base model with offloading support.
+base_model = AutoModelForCausalLM.from_pretrained(
+    model_dir,
+    device_map="auto",         # Automatically use available devices (GPU/CPU).
+    offload_folder=offload_dir # Directory to offload layers (for larger models).
+)
 
 # Load the adapter (PEFT) weights.
 model = PeftModel.from_pretrained(base_model, model_dir)
@@ -26,7 +28,7 @@ def generate_response(prompt, max_new_tokens=200, temperature=0.7):
     """
     # Tokenize the prompt and move tensors to the model's device.
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
+
     # Generate output text using sampling.
     outputs = model.generate(
         **inputs,
@@ -34,7 +36,7 @@ def generate_response(prompt, max_new_tokens=200, temperature=0.7):
         do_sample=True,
         temperature=temperature
     )
-    
+
     # Decode the generated tokens and return the response.
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response