enabled streaming

app.py

@@ -165,8 +165,63 @@ def chatbot_response(message, history):
     return response + metrics
 
 
+def process_dialog_streaming(message, history):
+    terminators = [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("<|eot_id|>"),
+    ]
+
+    dialog = [
+        {"role": "user" if i % 2 == 0 else "assistant", "content": msg}
+        for i, (msg, _) in enumerate(history)
+    ]
+    dialog.append({"role": "user", "content": message})
+
+    prompt = tokenizer.apply_chat_template(
+        dialog, tokenize=False, add_generation_prompt=True
+    )
+    tokenized_input_prompt_ids = tokenizer(
+        prompt, return_tensors="pt"
+    ).input_ids.to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    generation_kwargs = dict(
+        inputs=tokenized_input_prompt_ids,
+        streamer=streamer,
+        max_new_tokens=512,
+        temperature=0.4,
+        do_sample=True,
+        eos_token_id=terminators,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+
+    start_time = time.time()
+    total_tokens = 0
+
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    generated_text = ""
+    for new_text in streamer:
+        generated_text += new_text
+        total_tokens += 1
+        current_time = time.time()
+        elapsed_time = current_time - start_time
+        tokens_per_second = total_tokens / elapsed_time if elapsed_time > 0 else 0
+        print(f"Tokens per second: {tokens_per_second:.2f}", end="\r")
+        yield generated_text, elapsed_time, tokens_per_second
+
+    thread.join()
+
+def chatbot_response_streaming(message, history):
+    for response, generation_time, tokens_per_second in process_dialog_streaming(message, history):
+        metrics = f"\n\n---\n\n **Metrics**\t*Answer Generation Time:* `{generation_time:.2f} sec`\t*Tokens per Second:* `{tokens_per_second:.2f}`\n\n"
+        yield response + metrics
+
+
 demo = gr.ChatInterface(
-    fn=chatbot_response,
+    fn=chatbot_response_streaming,
     examples=["Hello", "How are you?", "Tell me a joke"],
     title="Chat with xMAD's: 1-bit-Llama-3-8B-Instruct Model",
     description="Contact support@xmad.ai to set up a demo",