Update app.py

app.py

@@ -3,20 +3,32 @@ import spaces
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
+from peft import PeftModel
 
-tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/BhashiniLLM")
-quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4",  
-        bnb_4bit_compute_dtype=torch.float16)
-model = AutoModelForCausalLM.from_pretrained("FlawedLLM/BhashiniLLM",
-                                             device_map="auto",
-                                             quantization_config=quantization_config, 
-                                             torch_dtype =torch.float16, 
-                                             low_cpu_mem_usage=True, 
-                                             use_safetensors=True,
-                                            )
+
+# tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/BhashiniLLM")
+# quantization_config = BitsAndBytesConfig(
+#         load_in_4bit=True,
+#         bnb_4bit_use_double_quant=True,
+#         bnb_4bit_quant_type="nf4",  
+#         bnb_4bit_compute_dtype=torch.float16)
+# model = AutoModelForCausalLM.from_pretrained("FlawedLLM/BhashiniLLM",
+#                                              device_map="auto",
+#                                              quantization_config=quantization_config, 
+#                                              torch_dtype =torch.float16, 
+#                                              low_cpu_mem_usage=True, 
+#                                              use_safetensors=True,
+#                                             )
+
+# Assuming you have your HF repository in this format: "your_username/your_model_name"
+model_id = "FlawedLLM/BhashiniLLM"
+
+# Load the base model (the one you fine-tuned with LoRA)
+base_model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')  # Load in 8-bit for efficiency
+
+# Load the LoRA adapter weights
+model = PeftModel.from_pretrained(base_model, model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
 @spaces.GPU(duration=300)