a

app.py

@@ -5,7 +5,7 @@ from huggingface_hub import InferenceClient
 Copied from inference in colab notebook
 """
 
-from transformers import AutoTokenizer , AutoModelForCausalLM , TextIteratorStreamer
+from transformers import AutoTokenizer , AutoModelForCausalLM , TextStreamer
 import torch
 from threading import Thread
 
@@ -93,22 +93,12 @@ def respond(
             messages.append({"role": "assistant", "content": val[1]})
     messages.append({"role": "user", "content": message})
     
-    # Create a single text prompt from the messages
-    prompt = ""
-    for msg in messages:
-        if msg["role"] == "system":
-            prompt += f"[System]: {msg['content']}\n\n"
-        elif msg["role"] == "user":
-            prompt += f"[User]: {msg['content']}\n\n"
-        elif msg["role"] == "assistant":
-            prompt += f"[Assistant]: {msg['content']}\n\n"
-    
     # Tokenize the prompt
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    inputs = tokenizer(messages, return_tensors="pt", truncation=True)
     input_ids = inputs.input_ids.to("cpu")  # Ensure input is on the CPU
 
     # Generate tokens incrementally
-    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    streamer = TextStreamer(tokenizer, skip_prompt=True)
     generation_kwargs = {
         "input_ids": input_ids,
         "max_new_tokens": max_tokens,
@@ -124,8 +114,7 @@ def respond(
     response = ""
     for token in streamer:
         response += token
-        yield response[len(prompt):].strip()
-        print(response)
+        yield response
 
 
 """