Update app.py

app.py

@@ -9,7 +9,7 @@ import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 import gradio as gr
 from threading import Thread
 
@@ -39,14 +39,18 @@ h3 {
 
 device = "cuda" # for GPU usage or "cpu" for CPU usage
 
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type= "nf4")
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     torch_dtype=torch.float16,
     device_map="auto",
-    load_in_8bit=False,
-    load_in_4bit=True,
-    use_flash_attention_2=True)
+    quantization_config=quantization_config)
 
 # Ensure `pad_token_id` is set
 if tokenizer.pad_token_id is None: