Update app.py

app.py

@@ -1,15 +1,8 @@
-# import subprocess
-# subprocess.run(
-#     'pip install flash-attn --no-build-isolation',
-#     env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
-#     shell=True
-# )
-
 import os
 import time
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 from threading import Thread
 
@@ -40,24 +33,12 @@ 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, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
     device_map="auto",
-    trust_remote_code=True,
-    # attn_implementation="flash_attention_2",
-    quantization_config=quantization_config).eval().to(device)
-
-# Ensure `pad_token_id` is set
-# if tokenizer.pad_token_id is None:
-#     tokenizer.pad_token_id = tokenizer.eos_token_id
+    trust_remote_code=True).eval().to(device)
 
 @spaces.GPU()
 def stream_chat(