Update app.py

app.py

@@ -1,10 +1,7 @@
 import os
-import time
-import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
-from threading import Thread
 
 MODEL_LIST = ["nawhgnuj/KamalaHarris-Llama-3.1-8B-Chat"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -60,8 +57,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
     quantization_config=quantization_config)
 
-@spaces.GPU()
-def stream_chat(
+def generate_response(
     message: str, 
     history: list,
     temperature: float,
@@ -91,33 +87,23 @@ Crucially, Keep responses concise and impactful."""
     conversation.append({"role": "user", "content": message})
     
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
-    attention_mask = torch.ones_like(input_ids)
     
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        repetition_penalty=repetition_penalty,
-        no_repeat_ngram_size=no_repeat_ngram_size,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id,
-        streamer=streamer,
-    )
-
     with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
-        
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
+        output = model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)
+    return response.strip()
 
 def add_text(history, text):
     history = history + [(text, None)]
@@ -125,11 +111,9 @@ def add_text(history, text):
 
 def bot(history, temperature, max_new_tokens, top_p, top_k, repetition_penalty, no_repeat_ngram_size):
     user_message = history[-1][0]
-    bot_response = stream_chat(user_message, history[:-1], temperature, max_new_tokens, top_p, top_k, repetition_penalty, no_repeat_ngram_size)
-    history[-1][1] = ""
-    for character in bot_response:
-        history[-1][1] += character
-        yield history
+    bot_response = generate_response(user_message, history[:-1], temperature, max_new_tokens, top_p, top_k, repetition_penalty, no_repeat_ngram_size)
+    history[-1][1] = bot_response
+    return history
 
 with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
     gr.HTML(TITLE)