Update app.py

app.py

@@ -5,20 +5,21 @@ import re
 import spaces
 import gradio as gr
 import torch
-# from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/Bhashini_gemma_merged4bit_clean_final")
-# model = AutoModelForCausalLM.from_pretrained("FlawedLLM/Bhashini_gemma_merged4bit_clean_final", load_in_4bit=True, device_map="auto")
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftModel
-
-# 1. Load Your Base Model and LoRA Adapter
-model_name_or_path = "FlawedLLM/Bhashini_gemma_merged4bit_clean_final"  # Hugging Face model or local path
-lora_weights = "FlawedLLM/Bhashini_gemma_lora_clean_final"  # LoRA weights
-
-tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
-model = AutoModelForCausalLM.from_pretrained(model_name_or_path, load_in_8bit=True, device_map='auto')
-model = PeftModel.from_pretrained(model, lora_weights)
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from huggingface_hub import login, HfFolder
+tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/Bhashini_gemma_merged16bit_clean_final", trust_remote_code=True)
+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/Bhashini_gemma_merged16bit_clean_final", 
+                                             device_map="auto",
+                                             quantization_config=quantization_config, 
+                                             torch_dtype =torch.float16, 
+                                             low_cpu_mem_usage=True, 
+                                             use_safetensors=True,
+                                             trust_remote_code=True)
 
 # alpaca_prompt = You MUST copy from above!
 @spaces.GPU(duration=300)