Streaming Assistant

README.md

@@ -224,7 +224,7 @@ chat_history = [
 num_tokens = asst.chat_completion_tokens(chat_history) #31
 print(num_tokens)
 
-result = asst.chat_completion(chat_history) #{'role': 'assistant', 'content': '1. Make a list of things to be grateful for.\n2. Go outside and take a walk in nature.\n3. Practice mindfulness meditation.\n4. Connect with a loved one or friend.\n5. Do something kind for someone else.\n6. Engage in a creative activity like drawing or writing.\n7. Read an uplifting book or listen to motivational podcasts.'}
+result = asst.chat_completion(chat_history, stream=True) #{'role': 'assistant', 'content': '1. Make a list of things to be grateful for.\n2. Go outside and take a walk in nature.\n3. Practice mindfulness meditation.\n4. Connect with a loved one or friend.\n5. Do something kind for someone else.\n6. Engage in a creative activity like drawing or writing.\n7. Read an uplifting book or listen to motivational podcasts.'}
 print(result["content"])
 ```
 - To use a dynamic knowledge base, see [![Q&A Assistant Demo](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1hkNgpSmmUA-mzUibqz25xq-E8KYOLuVx?usp=sharing)

demos/Text2Text_LLM.ipynb

@@ -2416,7 +2416,7 @@
   "cells": [
     {
       "cell_type": "code",
-      "execution_count": 1,
+      "execution_count": null,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -2562,7 +2562,7 @@
         "id": "NL3SY2o5gwNh",
         "outputId": "f083d548-cad3-46f4-867b-d36843237b6a"
       },
-      "execution_count": 1,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
@@ -2703,7 +2703,7 @@
         },
         "outputId": "e87035e8-94ef-424e-d511-558aae1f371b"
       },
-      "execution_count": 4,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
@@ -2749,28 +2749,30 @@
         "num_tokens = asst.chat_completion_tokens(chat_history) #31\n",
         "print(num_tokens)\n",
         "\n",
-        "result = asst.chat_completion(chat_history) #{'role': 'assistant', 'content': '1. Make a list of things to be grateful for.\\n2. Go outside and take a walk in nature.\\n3. Practice mindfulness meditation.\\n4. Connect with a loved one or friend.\\n5. Do something kind for someone else.\\n6. Engage in a creative activity like drawing or writing.\\n7. Read an uplifting book or listen to motivational podcasts.'}\n",
+        "result = asst.chat_completion(chat_history, stream=True) #{'role': 'assistant', 'content': '1. Make a list of things to be grateful for.\\n2. Go outside and take a walk in nature.\\n3. Practice mindfulness meditation.\\n4. Connect with a loved one or friend.\\n5. Do something kind for someone else.\\n6. Engage in a creative activity like drawing or writing.\\n7. Read an uplifting book or listen to motivational podcasts.'}\n",
         "print(result[\"content\"])"
       ],
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
         },
         "id": "TWONAnBZeoSO",
-        "outputId": "207c8783-3f24-43bc-81ce-d5b9ff296c28"
+        "outputId": "8ca0e813-cbdb-4d65-fc0f-4691f2c06364"
       },
       "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",
           "name": "stdout",
           "text": [
-            "Chat Completion API\n",
             "31\n",
-            "1. Start by making a to-do list for the day.\n",
-            "2. Prioritize tasks based on their importance and deadline.\n",
-            "3. Take breaks throughout the day to avoid burnout.\n",
-            "4. Make time for self-care activities like exercise or meditation.\n"
+            "1. Start your day with a healthy breakfast to fuel your body and mind.\n",
+            "2. Make time for some form of physical activity, such as going for a walk or run, practicing yoga, or joining a fitness class.\n",
+            "3. Set small, achievable goals for the day and work towards them one step at a time.\n",
+            "4. Take breaks throughout the day to recharge and refocus.\n",
+            "5. Connect with friends, family, or loved ones through phone calls, video chats, or in-person visits if possible.\n",
+            "6. Find something that brings you joy and make time for it in your schedule.\n",
+            "7. Practice gratitude by reflecting on the positive aspects of your life.\n"
           ]
         }
       ]

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
   name="text2text",
-  version="1.3.7",
+  version="1.3.8",
   author="artitw",
   author_email="artitw@gmail.com",
   description="Text2Text: Crosslingual NLP/G toolkit",

text2text/assistant.py

@@ -1,7 +1,7 @@
 import logging
 import pandas as pd
 import text2text as t2t
-from transformers import AutoModelForCausalLM, AutoTokenizer, logging
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, logging
 
 logging.set_verbosity(logging.CRITICAL)
 
@@ -86,13 +86,16 @@ def chat_completion(self, messages, **kwargs):
     top_k = kwargs.get('top_k', 0)
     repetition_penalty = kwargs.get('repetition_penalty', 1.15)
     max_new_tokens = kwargs.get('max_new_tokens', 512)
+    stream = kwargs.get('stream', False)
     tok = self.__class__.tokenizer
     m = self.__class__.model
+    streamer = TextStreamer(tok, skip_prompt=True, skip_special_tokens=True) if stream else None
 
     input_ids = tok([input_prompt], return_tensors="pt", padding=True).input_ids
     input_ids = input_ids.to(m.device)
     generate_kwargs = dict(
         input_ids=input_ids,
+        streamer=streamer,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         do_sample=temperature > 0.0,