From 84bb4b17c9d70785f6e0bbe74eaff4d3a84e4cc7 Mon Sep 17 00:00:00 2001 From: ascender1729 Date: Tue, 13 May 2025 20:08:34 +0530 Subject: [PATCH] feat: enhance LiteLLM tokenizer with flexible model support and improved error handling --- swarms/utils/litellm_tokenizer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/swarms/utils/litellm_tokenizer.py b/swarms/utils/litellm_tokenizer.py index c2743b10..3655bfdb 100644 --- a/swarms/utils/litellm_tokenizer.py +++ b/swarms/utils/litellm_tokenizer.py @@ -16,5 +16,13 @@ def count_tokens(text: str, model: str = "gpt-4o") -> int: return len(encode(model=model, text=text)) +class LiteLLMTokenizer: + def __init__(self, model_name="gpt-4o"): + self.model_name = model_name + + def count(self, text): + return count_tokens(text, model=self.model_name) + + # if __name__ == "__main__": # print(count_tokens("Hello, how are you?"))