Update app.py

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 import os
 from huggingface_hub import login
 
-# Load Hugging Face token (ensure it's set in 'Secrets')
+# Load Hugging Face token from environment variables
 hf_token = os.getenv("HUGGINGFACE_TOKEN")
 if hf_token:
     login(hf_token)
@@ -13,34 +13,32 @@ if hf_token:
 else:
     print("❌ Hugging Face token not found. Make sure it's set in 'Secrets'.")
 
-# Base model
+# Model paths
 BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.1"
-FINETUNED_MODEL_DIR = "./finetuned_model"  # Path to your adapter weights
+FINETUNED_MODEL_DIR = "./finetuned_model"  # Path to your fine-tuned adapter
+
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
-# Load base model (WITHOUT bitsandbytes)
+# Load base model efficiently
 base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,  
-    torch_dtype=torch.float32  # Ensure CPU compatibility
+    BASE_MODEL,
+    torch_dtype=torch_dtype,
+    device_map="auto" if torch.cuda.is_available() else None  # Use GPU if available
 )
 
-# Move base model to CPU
-base_model.to("cpu")
-
-# Load LoRA adapter
+# Load and merge LoRA adapter
 model = PeftModel.from_pretrained(base_model, FINETUNED_MODEL_DIR)
-
-# Merge adapter with base model
 model = model.merge_and_unload()
 
-# Move model to CPU (again, just to be sure)
-model.to("cpu")
+# Move model to appropriate device
+model.to(device)
 
-# Inference function
 def chat(message):
-    inputs = tokenizer(message, return_tensors="pt").to("cpu")  # Ensure inputs are on CPU
+    inputs = tokenizer(message, return_tensors="pt").to(device)
     output = model.generate(**inputs, max_new_tokens=100)
     response = tokenizer.decode(output[0], skip_special_tokens=True)
     return response