Update app.py

app.py

@@ -4,11 +4,15 @@ import torch
 import sqlite3
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
+import torch
+import os
+os.makedirs("offload", exist_ok=True)
 
 # ✅ Load fine-tuned models from Hugging Face Model Hub instead of Kaggle paths
 codellama_model_path = "srishtirai/codellama-sql-finetuned"  # Upload to HF Model Hub
 mistral_model_path = "srishtirai/mistral-sql-finetuned"  # Upload to HF Model Hub
 
+
 def load_model(model_path):
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     tokenizer.pad_token = tokenizer.eos_token
@@ -16,15 +20,19 @@ def load_model(model_path):
 
     peft_config = PeftConfig.from_pretrained(model_path)
     base_model_name = peft_config.base_model_name_or_path
+
     base_model = AutoModelForCausalLM.from_pretrained(
         base_model_name,
-        torch_dtype=torch.float16,
-        device_map="auto"
+        torch_dtype=torch.float16,  # Use FP16 to save memory
+        device_map="auto",  # Automatically allocate layers to CPU/GPU
+        offload_folder="offload"  # ✅ Offload large layers to disk
     )
+
     model = PeftModel.from_pretrained(base_model, model_path)
     model.eval()
     return model, tokenizer
 
+
 # ✅ Load both models from Hugging Face
 codellama_model, codellama_tokenizer = load_model(codellama_model_path)
 mistral_model, mistral_tokenizer = load_model(mistral_model_path)