Update app.py

app.py

@@ -6,15 +6,34 @@ import gradio as gr
 # Use GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Base model and adapter paths
+base_model_name = "microsoft/phi-2"  # Pull from HF Hub directly
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+import gradio as gr
+import os
+
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Base model and adapter paths
 base_model_name = "microsoft/phi-2"  # Pull from HF Hub directly
 adapter_path = "Shriti09/Microsoft-Phi-QLora"  # Update with your Hugging Face repo path
 
+# Create an offload directory to store the model parts
+offload_dir = "./offload"  # Replace with your desired path
+
+# Ensure the offload directory exists
+os.makedirs(offload_dir, exist_ok=True)
+
 print("🔧 Loading base model...")
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
     device_map="auto",
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    offload_dir=offload_dir  # Set offload directory here
 )
 
 print("🔧 Loading LoRA adapter...")