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/DonaldTrump-Llama-3.1-8B-Chat"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -59,14 +56,17 @@ 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,
+    max_new_tokens: int,
+    top_p: float,
+    top_k: int,
 ):
     system_prompt = """You are a Donald Trump chatbot. You only answer like Trump in his style and tone, reflecting his unique speech patterns. Incorporate the following characteristics in every response: 
     1. repeat key phrases for emphasis, use strong superlatives like 'tremendous' and 'fantastic,' attack opponents where appropriate (e.g., 'fake news media,' 'radical left')
-    2. focus on personal successes ('nobody\u2019s done more than I have')
+    2. focus on personal successes ('nobody's done more than I have')
     3. keep sentences short and impactful, and show national pride. 
     4. Maintain a direct, informal tone, often addressing the audience as 'folks' and dismiss opposing views bluntly. 
     5. Repeat key phrases for emphasis, but avoid excessive repetition.
@@ -84,19 +84,15 @@ 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)
-    attention_mask = torch.ones_like(input_ids)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
     
     with torch.no_grad():
         output = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=1024,
+            input_ids,
+            max_new_tokens=max_new_tokens,
             do_sample=True,
-            top_p=1.0,
-            top_k=20,
-            temperature=0.8,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
         )
@@ -108,13 +104,11 @@ def add_text(history, text):
     history = history + [(text, None)]
     return history, ""
 
-def bot(history):
+def bot(history, temperature, max_new_tokens, top_p, top_k):
     user_message = history[-1][0]
-    bot_response = stream_chat(user_message, history[:-1])
-    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)
+    history[-1][1] = bot_response
+    return history
 
 with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
     gr.HTML(TITLE)
@@ -135,6 +129,12 @@ with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
         submit = gr.Button("Submit", scale=1, variant="primary")
         clear = gr.Button("Clear", scale=1)
 
+    with gr.Accordion("Advanced Settings", open=False):
+        temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.8, step=0.1, label="Temperature")
+        max_new_tokens = gr.Slider(minimum=50, maximum=1024, value=1024, step=1, label="Max New Tokens")
+        top_p = gr.Slider(minimum=0.1, maximum=1.2, value=1.0, step=0.1, label="Top-p")
+        top_k = gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Top-k")
+
     gr.Examples(
         examples=[
             ["What's your stance on immigration?"],
@@ -145,11 +145,11 @@ with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
     )
 
     submit.click(add_text, [chatbot, msg], [chatbot, msg], queue=False).then(
-        bot, chatbot, chatbot
+        bot, [chatbot, temperature, max_new_tokens, top_p, top_k], chatbot
     )
     clear.click(lambda: [], outputs=[chatbot], queue=False)
     msg.submit(add_text, [chatbot, msg], [chatbot, msg], queue=False).then(
-        bot, chatbot, chatbot
+        bot, [chatbot, temperature, max_new_tokens, top_p, top_k], chatbot
     )
 
 if __name__ == "__main__":