Schema for structured outputs

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

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"

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": []
     }
   ]
 }

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',

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)