Update bitsandbytes configuration and model loading

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import BitsAndBytesConfig
 from accelerate import infer_auto_device_map
 
 # Load the model name
@@ -8,8 +9,17 @@ model_name = "ai4bharat/Airavata"
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Load the model first
-model = AutoModelForCausalLM.from_pretrained(model_name, load_in_8bit=True)
+# Create a BitsAndBytesConfig for quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True,  # Set this to True for 8-bit loading
+    # Optionally, you can specify more parameters based on your needs
+)
+
+# Load the model using the BitsAndBytesConfig
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config  # Use the BitsAndBytesConfig
+)
 
 # Now infer the device map
 device_map = infer_auto_device_map(model)