Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
 import torch
 
@@ -7,36 +7,31 @@ model_id = "saikrishnagorijala/friday-V1"
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-# Load model in 8-bit mode (requires bitsandbytes)
+# Define quantization config for 8-bit inference
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    bnb_8bit_use_double_quant=True,
+    bnb_8bit_quant_type="nf4",
+    bnb_8bit_compute_dtype=torch.float16
+)
+
+# Load model with quantization_config
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",   # automatically put layers on GPU
-    load_in_8bit=True,   # enable 8-bit quantization
-    torch_dtype=torch.float16  # keep computations in FP16 where needed
+    device_map="auto",
+    quantization_config=bnb_config
 )
 
 def chat(prompt):
-    # Tokenize input and move to model device
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    # Generate response
     outputs = model.generate(
         **inputs,
         max_new_tokens=200,
-        do_sample=True,       # allow sampling for varied responses
-        temperature=1.2,      # optional creativity control
+        do_sample=True,
+        temperature=1.2,
         top_p=0.9
     )
-
-    # Decode and return text
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Gradio interface
-demo = gr.Interface(
-    fn=chat,
-    inputs="text",
-    outputs="text",
-    title="Friday-V1 Chatbot"
-)
-
+demo = gr.Interface(fn=chat, inputs="text", outputs="text", title="Friday-V1 Chatbot")
 demo.launch()