Update app.py

app.py

@@ -6,13 +6,15 @@ from huggingface_hub import InferenceClient
 # """
 # client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
-from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM#, MambaForCausalLM
+from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM, TextStreamer
 from peft import PeftConfig, PeftModel
 
 config = PeftConfig.from_pretrained("jonathanjordan21/mos-mamba-6x130m-trainer")
 
 tokenizer = AutoTokenizer.from_pretrained("jonathanjordan21/mos-mamba-6x130m-trainer", trust_remote_code=True)
 
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
 model = AutoModelForCausalLM.from_pretrained(
     "jonathanjordan21/mos-mamba-6x130m-trainer",
     eos_token_id=tokenizer.eos_token_id,
@@ -25,6 +27,7 @@ model = model.merge_and_unload()
 print(model.config.eos_token_id)
 
 
+
 def invoke(
     message,
     history: list[tuple[str, str]],
@@ -45,11 +48,21 @@ def invoke(
 
     tokens = tokenizer.apply_chat_template(messages, return_tensors='pt', add_generation_prompt=True)
 
-    out = model.generate(tokens, eos_token_id=model.config.eos_token_id, max_new_tokens=max_tokens, repetition_penalty=1.05, temperature=temperature, top_p=top_p)
+    response = 
+
+    for res in model.generate(
+        tokens, 
+        streamer=streamer, 
+        eos_token_id=model.config.eos_token_id, 
+        max_new_tokens=max_tokens, 
+        temperature=temperature
+    ):
+        response += res
+        yield response
 
-    res = tokenizer.batch_decode(out)
+    # res = tokenizer.batch_decode(out)
 
-    yield res
+    # yield res
 
 
 def respond(