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, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import gradio as gr
 from threading import Thread
 
@@ -24,7 +17,6 @@ PLACEHOLDER = """
 </center>
 """
 
-
 CSS = """
 .duplicate-button {
     margin: auto !important;
@@ -37,27 +29,20 @@ 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")
+device = "cpu"  # Changed to CPU
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float32,  # Changed to float32 for better CPU compatibility
     device_map="auto",
-    attn_implementation="flash_attention_2",
-    quantization_config=quantization_config)
+).to(device)  # Explicitly move to CPU
 
 # Ensure `pad_token_id` is set
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
-@spaces.GPU()
+@spaces.CPU()  # Changed to CPU decorator
 def stream_chat(
     message: str, 
     history: list,
@@ -82,19 +67,19 @@ def stream_chat(
 
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
+    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(device)
     
     streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
     
     generate_kwargs = dict(
         input_ids=input_ids, 
-        max_new_tokens = max_new_tokens,
-        do_sample = False if temperature == 0 else True,
-        top_p = top_p,
-        top_k = top_k,
-        eos_token_id = tokenizer.eos_token_id,
-        pad_token_id = tokenizer.pad_token_id,
-        temperature = temperature,
+        max_new_tokens=max_new_tokens,
+        do_sample=False if temperature == 0 else True,
+        top_p=top_p,
+        top_k=top_k,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        temperature=temperature,
         repetition_penalty=penalty,
         streamer=streamer,
     )
@@ -107,7 +92,6 @@ def stream_chat(
     for new_text in streamer:
         buffer += new_text
         yield buffer
-
             
 chatbot = gr.Chatbot(height=600, placeholder=PLACEHOLDER)
 
@@ -138,7 +122,7 @@ with gr.Blocks(css=CSS, theme="soft") as demo:
                 minimum=128,
                 maximum=8192,
                 step=1,
-                value= 2500,
+                value=2500,
                 label="Max new tokens",
                 render=False,
             ),
@@ -176,6 +160,5 @@ with gr.Blocks(css=CSS, theme="soft") as demo:
         cache_examples=False,
     )
 
-
 if __name__ == "__main__":
     demo.launch()