From ba7877c1d3d0ae903e83bc7b0c946913fac4ca09 Mon Sep 17 00:00:00 2001 From: artitw Date: Sun, 4 Feb 2024 19:49:34 +0000 Subject: [PATCH] Streaming Assistant --- README.md | 2 +- demos/Text2Text_LLM.ipynb | 22 ++++++++++++---------- setup.py | 2 +- text2text/assistant.py | 5 ++++- 4 files changed, 18 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 1d7494c..27166e4 100755 --- a/README.md +++ b/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) diff --git a/demos/Text2Text_LLM.ipynb b/demos/Text2Text_LLM.ipynb index f73f83e..8fafb87 100644 --- a/demos/Text2Text_LLM.ipynb +++ b/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,7 +2749,7 @@ "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": { @@ -2757,7 +2757,7 @@ "base_uri": "https://localhost:8080/" }, "id": "TWONAnBZeoSO", - "outputId": "207c8783-3f24-43bc-81ce-d5b9ff296c28" + "outputId": "8ca0e813-cbdb-4d65-fc0f-4691f2c06364" }, "execution_count": null, "outputs": [ @@ -2765,12 +2765,14 @@ "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" ] } ] diff --git a/setup.py b/setup.py index 5e96ea9..2160beb 100755 --- a/setup.py +++ b/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", diff --git a/text2text/assistant.py b/text2text/assistant.py index 9a26a1b..0d2e549 100644 --- a/text2text/assistant.py +++ b/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,