Skip to content

Commit

Permalink
remove lambda
Browse files Browse the repository at this point in the history
  • Loading branch information
ILSparkle committed Jun 27, 2024
1 parent e21368a commit 22cb891
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions src/cardinal/model/token_counter.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,14 +14,15 @@
class TokenCounter:
def __init__(self, model: Optional[str] = None) -> None:
if settings.hf_tokenizer_path is not None:
tokenizer = AutoTokenizer.from_pretrained(
self._encoding = AutoTokenizer.from_pretrained(
settings.hf_tokenizer_path,
trust_remote_code=True,
)
self._encode_func = lambda text: len(tokenizer.tokenize(text))
else:
encoding = tiktoken.encoding_for_model(model if model is not None else settings.default_chat_model)
self._encode_func = lambda text: len(encoding.encode(text))
self._encoding = tiktoken.encoding_for_model(model if model is not None else settings.default_chat_model)

def __call__(self, text: str) -> int:
return self._encode_func(text)
if settings.hf_tokenizer_path is not None:
return len(self._encoding.tokenize(text))
else:
return len(self._encoding.encode(text))

0 comments on commit 22cb891

Please sign in to comment.