Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig, AutoConfig
 import torch
 
 MODEL_URL = "kingabzpro/Llama-3.1-8B-Instruct-Mental-Health-Classification"
@@ -8,8 +8,18 @@ MODEL_URL = "kingabzpro/Llama-3.1-8B-Instruct-Mental-Health-Classification"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
 tokenizer.pad_token_id = tokenizer.eos_token_id
 
-model = AutoModelForCausalLM.from_pretrained(MODEL_URL, low_cpu_mem_usage=True, return_dict=True,torch_dtype=torch.float16,
-    device_map="cpu")
+bnbConfig = BitsAndBytesConfig(
+    load_in_4bit = True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+model = AutoModelForCausalLM.from_pretrained(MODEL_URL, 
+                                             quantization_config=bnbConfig,
+                                             low_cpu_mem_usage=True, 
+                                             return_dict=True,
+                                             torch_dtype=torch.float16,
+                                             device_map="cpu")
 
 def prediction(text):
     # create pipeline