stream output and read from history

this will stream the output while generating, it will also make use of the history meaning the model will read the past conversation, i'm using the code from my other repo at : https://huggingface.co/spaces/not-lain/text-streaming/blob/main/app.py.

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import os
 import spaces
-from transformers import GemmaTokenizer, AutoModelForCausalLM
+from transformers import GemmaTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
 
 # Set an environment variable
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -24,14 +25,35 @@ def codegemma(message: str, history: list, temperature: float, max_new_tokens: i
     Returns:
         str: The generated response.
     """
-    input_ids = tokenizer(message, return_tensors="pt").to("cuda:0")
-    outputs = model.generate(
-        **input_ids,
-        temperature=temperature,
+    chat = []
+    for item in history:
+        chat.append({"role": "user", "content": item[0]})
+        if item[1] is not None:
+            chat.append({"role": "assistant", "content": item[1]})
+    chat.append({"role": "user", "content": message})
+    messages = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+    # Tokenize the messages string
+    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(
+        tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+    
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
         max_new_tokens=max_new_tokens,
+        temperature=temperature,
     )
-    response = tokenizer.decode(outputs[0])
-    return response
+    
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    # Initialize an empty string to store the generated text
+    partial_text = ""
+    for new_text in streamer:
+        # print(new_text)
+        partial_text += new_text
+        # Yield an empty string to cleanup the message textbox and the updated conversation history
+        yield partial_text
 
 
 placeholder = """