Update app.py

app.py

@@ -1,13 +1,22 @@
 import re
 import spaces
 import gradio as gr
-from peft import PeftModel, PeftConfig
-from peft import AutoPeftModelForCausalLM
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
 
 tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/BhashiniLLM")
-model = AutoModelForCausalLM.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,
+                                            )
 
 
 @spaces.GPU(duration=300)