Merge pull request #1556 from BerriAI/litellm_importlib_issue

fix(utils.py): move from pkg_resources to importlib
This commit is contained in:
Krish Dholakia 2024-01-22 15:56:07 -08:00 committed by GitHub
commit b1cced16fc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 11 deletions

View file

@ -2505,24 +2505,20 @@ def get_replicate_completion_pricing(completion_response=None, total_time=0.0):
def _select_tokenizer(model: str):
# cohere
import pkg_resources
from importlib import resources
if model in litellm.cohere_models:
# cohere
tokenizer = Tokenizer.from_pretrained("Cohere/command-nightly")
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
# anthropic
elif model in litellm.anthropic_models:
# Read the JSON file
filename = pkg_resources.resource_filename(
__name__, "llms/tokenizers/anthropic_tokenizer.json"
)
with open(filename, "r") as f:
with resources.open_text(
"litellm.llms.tokenizers", "anthropic_tokenizer.json"
) as f:
json_data = json.load(f)
# Decode the JSON data from utf-8
json_data_decoded = json.dumps(json_data, ensure_ascii=False)
# Convert to str
json_str = str(json_data_decoded)
# Convert to str (if necessary)
json_str = json.dumps(json_data)
# load tokenizer
tokenizer = Tokenizer.from_str(json_str)
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}