Add tests + cleanup

examples/README.md

@@ -36,6 +36,10 @@ python3 examples/<example>.py
 - [structured-outputs-image.py](structured-outputs-image.py)
 
 
+### Tokenization - Tokenize and detokenize text with a model
+- [tokenization.py](tokenization.py)
+
+
 ### Ollama List - List all downloaded models and their properties
 - [list.py](list.py)
 

examples/structured-outputs-image.py

@@ -2,14 +2,13 @@ from pathlib import Path
 from pydantic import BaseModel
 from typing import List, Optional, Literal
 from ollama import chat
-from rich import print
 
 
 # Define the schema for image objects
 class Object(BaseModel):
   name: str
   confidence: float
-  attributes: Optional[dict] = None
+  attributes: str
 
 
 class ImageDescription(BaseModel):

examples/tokenization.py

@@ -0,0 +1,10 @@
+import ollama
+
+# Get tokens from a model
+response = ollama.tokenize(model='llama3.2', text='Why the sky is blue?')
+tokens = response.tokens
+print('Tokens from model', tokens)
+
+# Convert tokens back to text
+response = ollama.detokenize(model='llama3.2', tokens=tokens)
+print('Text from tokens', response.text)  # Prints: Why the sky is blue?

ollama/__init__.py

@@ -12,6 +12,8 @@ from ollama._types import (
   ListResponse,
   ShowResponse,
   ProcessResponse,
+  TokenizeResponse,
+  DetokenizeResponse,
   RequestError,
   ResponseError,
 )
@@ -31,6 +33,8 @@ __all__ = [
   'ListResponse',
   'ShowResponse',
   'ProcessResponse',
+  'TokenizeResponse',
+  'DetokenizeResponse',
   'RequestError',
   'ResponseError',
 ]
@@ -49,3 +53,5 @@ list = _client.list
 copy = _client.copy
 show = _client.show
 ps = _client.ps
+tokenize = _client.tokenize
+detokenize = _client.detokenize

ollama/_client.py

@@ -48,6 +48,8 @@ from ollama._types import (
   CreateRequest,
   CopyRequest,
   DeleteRequest,
+  DetokenizeRequest,
+  DetokenizeResponse,
   EmbedRequest,
   EmbedResponse,
   EmbeddingsRequest,
@@ -67,6 +69,8 @@ from ollama._types import (
   ShowRequest,
   ShowResponse,
   StatusResponse,
+  TokenizeRequest,
+  TokenizeResponse,
   Tool,
 )
 
@@ -188,7 +192,7 @@ class Client(BaseClient):
     context: Optional[Sequence[int]] = None,
     stream: Literal[False] = False,
     raw: bool = False,
-    format: Optional[Union[Literal['json'], JsonSchemaValue]] = None,
+    format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
     images: Optional[Sequence[Union[str, bytes]]] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
     keep_alive: Optional[Union[float, str]] = None,
@@ -206,7 +210,7 @@ class Client(BaseClient):
     context: Optional[Sequence[int]] = None,
     stream: Literal[True] = True,
     raw: bool = False,
-    format: Optional[Union[Literal['json'], JsonSchemaValue]] = None,
+    format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
     images: Optional[Sequence[Union[str, bytes]]] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
     keep_alive: Optional[Union[float, str]] = None,
@@ -223,7 +227,7 @@ class Client(BaseClient):
     context: Optional[Sequence[int]] = None,
     stream: bool = False,
     raw: Optional[bool] = None,
-    format: Optional[Union[Literal['json'], JsonSchemaValue]] = None,
+    format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
     images: Optional[Sequence[Union[str, bytes]]] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
     keep_alive: Optional[Union[float, str]] = None,
@@ -267,7 +271,7 @@ class Client(BaseClient):
     *,
     tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
     stream: Literal[False] = False,
-    format: Optional[Union[Literal['json'], JsonSchemaValue]] = None,
+    format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
     keep_alive: Optional[Union[float, str]] = None,
   ) -> ChatResponse: ...
@@ -280,7 +284,7 @@ class Client(BaseClient):
     *,
     tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
     stream: Literal[True] = True,
-    format: Optional[Union[Literal['json'], JsonSchemaValue]] = None,
+    format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
     keep_alive: Optional[Union[float, str]] = None,
   ) -> Iterator[ChatResponse]: ...
@@ -292,7 +296,7 @@ class Client(BaseClient):
     *,
     tools: Optional[Sequence[Union[Mapping[str, Any], Tool, Callable]]] = None,
     stream: bool = False,
-    format: Optional[Union[Literal['json'], JsonSchemaValue]] = None,
+    format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None,
     options: Optional[Union[Mapping[str, Any], Options]] = None,
     keep_alive: Optional[Union[float, str]] = None,
   ) -> Union[ChatResponse, Iterator[ChatResponse]]:
@@ -611,6 +615,28 @@ class Client(BaseClient):
       '/api/ps',
     )
 
+  def tokenize(self, model: str, text: str) -> TokenizeResponse:
+    return self._request(
+      TokenizeResponse,
+      'POST',
+      '/api/tokenize',
+      json=TokenizeRequest(
+        model=model,
+        text=text,
+      ).model_dump(exclude_none=True),
+    )
+
+  def detokenize(self, model: str, tokens: Sequence[int]) -> DetokenizeResponse:
+    return self._request(
+      DetokenizeResponse,
+      'POST',
+      '/api/detokenize',
+      json=DetokenizeRequest(
+        model=model,
+        tokens=tokens,
+      ).model_dump(exclude_none=True),
+    )
+
 
 class AsyncClient(BaseClient):
   def __init__(self, host: Optional[str] = None, **kwargs) -> None:
@@ -1120,6 +1146,28 @@ class AsyncClient(BaseClient):
       '/api/ps',
     )
 
+  async def tokenize(self, model: str, text: str) -> TokenizeResponse:
+    return await self._request(
+      TokenizeResponse,
+      'POST',
+      '/api/tokenize',
+      json=TokenizeRequest(
+        model=model,
+        text=text,
+      ).model_dump(exclude_none=True),
+    )
+
+  async def detokenize(self, model: str, tokens: Sequence[int]) -> DetokenizeResponse:
+    return await self._request(
+      DetokenizeResponse,
+      'POST',
+      '/api/detokenize',
+      json=DetokenizeRequest(
+        model=model,
+        tokens=tokens,
+      ).model_dump(exclude_none=True),
+    )
+
 
 def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]:
   for message in messages or []:

ollama/_types.py

@@ -151,7 +151,7 @@ class BaseGenerateRequest(BaseStreamableRequest):
   options: Optional[Union[Mapping[str, Any], Options]] = None
   'Options to use for the request.'
 
-  format: Optional[Union[Literal['json'], JsonSchemaValue]] = None
+  format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None
   'Format of the response.'
 
   keep_alive: Optional[Union[float, str]] = None
@@ -494,6 +494,24 @@ class ProcessResponse(SubscriptableBaseModel):
   models: Sequence[Model]
 
 
+class TokenizeRequest(BaseRequest):
+  model: str
+  text: str
+
+
+class TokenizeResponse(BaseGenerateResponse):
+  tokens: Sequence[int]
+
+
+class DetokenizeRequest(BaseRequest):
+  model: str
+  tokens: Sequence[int]
+
+
+class DetokenizeResponse(BaseGenerateResponse):
+  text: str
+
+
 class RequestError(Exception):
   """
   Common class for request errors.

tests/test_client.py

@@ -1260,3 +1260,27 @@ def test_tool_validation():
   with pytest.raises(ValidationError):
     invalid_tool = {'type': 'invalid_type', 'function': {'name': 'test'}}
     list(_copy_tools([invalid_tool]))
+
+
+def test_client_tokenize(httpserver: HTTPServer):
+  httpserver.expect_ordered_request(
+    '/api/tokenize',
+    method='POST',
+    json={'model': 'dummy', 'text': 'Hello world!'},
+  ).respond_with_json({'tokens': [1, 2, 3]})
+
+  client = Client(httpserver.url_for('/'))
+  response = client.tokenize('dummy', 'Hello world!')
+  assert response.tokens == [1, 2, 3]
+
+
+def test_client_detokenize(httpserver: HTTPServer):
+  httpserver.expect_ordered_request(
+    '/api/detokenize',
+    method='POST',
+    json={'model': 'dummy', 'tokens': [1, 2, 3]},
+  ).respond_with_json({'text': 'Hello world!'})
+
+  client = Client(httpserver.url_for('/'))
+  response = client.detokenize('dummy', [1, 2, 3])
+  assert response.text == 'Hello world!'