Update app.py

app.py

@@ -73,13 +73,6 @@ def stream_chat(
     Importantly, always respond to points in Trump's style. Keep responses concise and avoid unnecessary repetition.
     """
 
-    temperature = 0.1
-    max_new_tokens = 256
-    top_p = 0.9
-    top_k = 20
-    repetition_penalty = 1.5
-    no_repeat_ngram_size = 3
-    
     conversation = [
         {"role": "system", "content": system_prompt}
     ]
@@ -95,29 +88,21 @@ def stream_chat(
     
     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=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=1024,
+            do_sample=True,
+            top_p=1.0,
+            top_k=20,
+            temperature=0.8,
+            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)]