From 422724ee8c923f19df71e65727984d03b6e3dcca Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Sun, 15 Dec 2024 23:53:43 -0800
Subject: [PATCH] Add tokenize detokenize compatibility

---
 examples/tokenization.py | 10 +++++++++
 ollama/__init__.py       |  6 +++++
 ollama/_client.py        | 48 ++++++++++++++++++++++++++++++++++++++++
 ollama/_types.py         | 18 +++++++++++++++
 4 files changed, 82 insertions(+)
 create mode 100644 examples/tokenization.py

diff --git a/examples/tokenization.py b/examples/tokenization.py
new file mode 100644
index 0000000..68a033d
--- /dev/null
+++ b/examples/tokenization.py
@@ -0,0 +1,10 @@
+import ollama
+
+# Get tokens from a model
+response = ollama.tokenize(model='llama3.2', text='Hello world!')
+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: Hello world!
diff --git a/ollama/__init__.py b/ollama/__init__.py
index 23d736a..049da05 100644
--- a/ollama/__init__.py
+++ b/ollama/__init__.py
@@ -12,6 +12,8 @@
   ListResponse,
   ShowResponse,
   ProcessResponse,
+  TokenizeResponse,
+  DetokenizeResponse,
   RequestError,
   ResponseError,
 )
@@ -31,6 +33,8 @@
   'ListResponse',
   'ShowResponse',
   'ProcessResponse',
+  'TokenizeResponse',
+  'DetokenizeResponse',
   'RequestError',
   'ResponseError',
 ]
@@ -49,3 +53,5 @@
 copy = _client.copy
 show = _client.show
 ps = _client.ps
+tokenize = _client.tokenize
+detokenize = _client.detokenize
diff --git a/ollama/_client.py b/ollama/_client.py
index 87fa881..e2fc059 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -48,6 +48,8 @@
   CreateRequest,
   CopyRequest,
   DeleteRequest,
+  DetokenizeRequest,
+  DetokenizeResponse,
   EmbedRequest,
   EmbedResponse,
   EmbeddingsRequest,
@@ -67,6 +69,8 @@
   ShowRequest,
   ShowResponse,
   StatusResponse,
+  TokenizeRequest,
+  TokenizeResponse,
   Tool,
 )
 
@@ -611,6 +615,28 @@ def ps(self) -> ProcessResponse:
       '/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 @@ async def ps(self) -> ProcessResponse:
       '/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 []:
diff --git a/ollama/_types.py b/ollama/_types.py
index 3be80a7..aa6dfe8 100644
--- a/ollama/_types.py
+++ b/ollama/_types.py
@@ -494,6 +494,24 @@ class Model(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.