Update app.py

app.py

@@ -1,16 +1,18 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 
 model_id = "thrishala/mental_health_chatbot"
 
 try:
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,  # we going to 4 babey
+    )
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        load_in_8bit=True,  # Load in 8-bit quantization
+        quantization_config=quantization_config,
         device_map="auto", #Use GPU if available
-        torch_dtype=torch.float16 #Use float 16 for additional memory reduction
     )
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)