Add tokenize detokenize compatibility

This commit is contained in:
ParthSareen
2024-12-16 09:41:07 -08:00
parent 70dd0b7e63
commit 422724ee8c
4 changed files with 82 additions and 0 deletions
+10
View File
@@ -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!
+6
View File
@@ -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
+48
View File
@@ -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,
)
@@ -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 []:
+18
View File
@@ -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.