diff --git a/README.md b/README.md index a165c92..211cdf5 100755 --- a/README.md +++ b/README.md @@ -202,7 +202,32 @@ chat_history = [ 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.'} for chunk in result: - print(chunk['message']['content'], end='', flush=True) + print(chunk.raw['message']['content'], end='', flush=True) + +# Running conversation +messages = [] +while True: + user_input = input("User: ") + print() + messages.append({"role": "user", "content": user_input}) + print("Assistant: ") + result = asst.chat_completion(messages, stream=False) + print(result["message"]["content"]) + messages.append(result["message"]) + print() + +# Schema for structured output +from pydantic import BaseModel + +class Song(BaseModel): + name: str + artist: str + +result = asst.chat_completion([ + {"role": "user", "content": "What is Britney Spears's best song?"} +], schema=Song, max_new_tokens=16) + +# Song(name='Toxic', artist='Britney Spears') ``` ### Tokenization diff --git a/demos/Text2Text_Demos.ipynb b/demos/Text2Text_Demos.ipynb index 1fe54af..4a40878 100644 --- a/demos/Text2Text_Demos.ipynb +++ b/demos/Text2Text_Demos.ipynb @@ -62,6 +62,7 @@ "# Free private open source alternative to commercial LLMs.\n", "# Commercial LLMs are costly, collect your data, impose quotas and rate limits that hinder development.\n", "# Run at no cost on Google Colab free tier, so you don't even need your own device.\n", + "# For more examples, see https://colab.research.google.com/drive/1K6Kk80w9vjFZ7PL9dPRgVuOPuaWcY4ae\n", "# To add a knowledge base, see https://colab.research.google.com/drive/1hkNgpSmmUA-mzUibqz25xq-E8KYOLuVx?usp=sharing\n", "\n", "asst = t2t.Assistant()\n", @@ -75,7 +76,7 @@ "\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", "for chunk in result:\n", - " print(chunk['message']['content'], end='', flush=True)" + " print(chunk.raw['message']['content'], end='', flush=True)" ], "metadata": { "id": "VPMdUSy9YYRl" diff --git a/demos/Text2Text_LLM.ipynb b/demos/Text2Text_LLM.ipynb index 3c804d1..eeab08b 100644 --- a/demos/Text2Text_LLM.ipynb +++ b/demos/Text2Text_LLM.ipynb @@ -93,7 +93,7 @@ "\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", "for chunk in result:\n", - " print(chunk['message']['content'], end='', flush=True)" + " print(chunk.raw['message']['content'], end='', flush=True)" ], "metadata": { "id": "TWONAnBZeoSO" @@ -161,6 +161,26 @@ ] } ] + }, + { + "cell_type": "code", + "source": [ + "# Schema for structured output\n", + "from pydantic import BaseModel\n", + "\n", + "class Song(BaseModel):\n", + " name: str\n", + " artist: str\n", + "\n", + "result = asst.chat_completion([\n", + " {\"role\": \"user\", \"content\": \"What is Britney Spears's best song?\"}\n", + "], schema=Song, max_new_tokens=16)" + ], + "metadata": { + "id": "e5khHlNQZ0FD" + }, + "execution_count": null, + "outputs": [] } ] } \ No newline at end of file diff --git a/setup.py b/setup.py index 7f02187..5be065f 100755 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name="text2text", - version="1.4.7", + version="1.4.8", author="artitw", author_email="artitw@gmail.com", description="Text2Text: Crosslingual NLP/G toolkit", @@ -23,9 +23,10 @@ 'faiss-cpu', 'flask', 'googledrivedownloader', + 'llama-index-llms-ollama', 'numpy', - 'ollama', 'pandas', + 'pydantic', 'scikit-learn', 'scipy', 'sentencepiece', diff --git a/text2text/assistant.py b/text2text/assistant.py index 1f94664..7064cb3 100644 --- a/text2text/assistant.py +++ b/text2text/assistant.py @@ -1,6 +1,8 @@ import os import ollama import psutil +from llama_index.llms.ollama import Ollama +from llama_index.core.llms import ChatMessage def is_port_in_use(port): for conn in psutil.net_connections(): @@ -19,16 +21,22 @@ def __init__(self, **kwargs): if return_code != 0: print("Cannot install ollama.") self.load_model() - self.client = ollama.Client(host=self.model_url) + self.client = Ollama(model=self.model_name, request_timeout=120.0) def load_model(self): + return_code = os.system("sudo service ollama stop") return_code = os.system(f"ollama serve & ollama pull {self.model_name}") if return_code != 0: - print(f"{self.model_name} is not loading up. Maybe needs more memory. Restarting and trying again might help.") + print(f"{self.model_name} is not loading up. Restarting and trying again might help. Maybe needs more memory.") - def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, **kwargs): + def chat_completion(self, messages=[{"role": "user", "content": "hello"}], stream=False, schema=None, **kwargs): if is_port_in_use(self.port): - return self.client.chat(model=self.model_name, messages=messages, stream=stream) + msgs = [ChatMessage(**m) for m in messages] + if stream: + return self.client.stream_chat(messages=msgs) + if schema: + return self.client.as_structured_llm(schema).chat(messages=msgs).raw + return self.client.chat(messages=msgs).raw self.load_model() return self.chat_completion(messages=messages, stream=stream, **kwargs)