Update app.py

app.py

@@ -1,19 +1,12 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
-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
+from accelerate import Accelerator
+
+# Initialize the Accelerator to manage device placement and offloading
+accelerator = Accelerator()
 
 # Use GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -23,19 +16,22 @@ 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
+offload_dir = "./offload_dir"  # Set offload directory here
 
 # Ensure the offload directory exists
 os.makedirs(offload_dir, exist_ok=True)
 
 print("🔧 Loading base model...")
+# Using the Accelerator to load the model and dispatch to the correct devices
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
-    device_map="auto",
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    offload_dir=offload_dir  # Set offload directory here
+    device_map="auto",  # Automatically map the model to available devices
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
 )
 
+# Ensure the model is offloaded when necessary
+base_model = accelerator.prepare(base_model)
+
 print("🔧 Loading LoRA adapter...")
 adapter_model = PeftModel.from_pretrained(base_model, adapter_path)