stream output and read from history

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 = """