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, ListResponse,
ShowResponse, ShowResponse,
ProcessResponse, ProcessResponse,
TokenizeResponse,
DetokenizeResponse,
RequestError, RequestError,
ResponseError, ResponseError,
) )
@@ -31,6 +33,8 @@ __all__ = [
'ListResponse', 'ListResponse',
'ShowResponse', 'ShowResponse',
'ProcessResponse', 'ProcessResponse',
'TokenizeResponse',
'DetokenizeResponse',
'RequestError', 'RequestError',
'ResponseError', 'ResponseError',
] ]
@@ -49,3 +53,5 @@ list = _client.list
copy = _client.copy copy = _client.copy
show = _client.show show = _client.show
ps = _client.ps ps = _client.ps
tokenize = _client.tokenize
detokenize = _client.detokenize
+48
View File
@@ -48,6 +48,8 @@ from ollama._types import (
CreateRequest, CreateRequest,
CopyRequest, CopyRequest,
DeleteRequest, DeleteRequest,
DetokenizeRequest,
DetokenizeResponse,
EmbedRequest, EmbedRequest,
EmbedResponse, EmbedResponse,
EmbeddingsRequest, EmbeddingsRequest,
@@ -67,6 +69,8 @@ from ollama._types import (
ShowRequest, ShowRequest,
ShowResponse, ShowResponse,
StatusResponse, StatusResponse,
TokenizeRequest,
TokenizeResponse,
Tool, Tool,
) )
@@ -611,6 +615,28 @@ class Client(BaseClient):
'/api/ps', '/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): class AsyncClient(BaseClient):
def __init__(self, host: Optional[str] = None, **kwargs) -> None: def __init__(self, host: Optional[str] = None, **kwargs) -> None:
@@ -1120,6 +1146,28 @@ class AsyncClient(BaseClient):
'/api/ps', '/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]: def _copy_messages(messages: Optional[Sequence[Union[Mapping[str, Any], Message]]]) -> Iterator[Message]:
for message in messages or []: for message in messages or []:
+18
View File
@@ -494,6 +494,24 @@ class ProcessResponse(SubscriptableBaseModel):
models: Sequence[Model] 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): class RequestError(Exception):
""" """
Common class for request errors. Common class for request errors.