ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1495 from puffo/litellm_ollama_chat_fix

Browse Source 
fix(ollama_chat.py): use tiktoken as backup for prompt token counting
This commit is contained in:
Krish Dholakia 2024-01-18 10:02:27 -08:00 committed by GitHub
commit 658fd4de38
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
litellm/llms/ollama_chat.py
View File

@ -220,7 +220,7 @@ def get_ollama_response(
model_response["choices"][0]["message"] = response_json["message"]
model_response["created"] = int(time.time())
model_response["model"] = "ollama/" + model
prompt_tokens = response_json["prompt_eval_count"] # type: ignore
prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore
completion_tokens = response_json["eval_count"]
model_response["usage"] = litellm.Usage(
prompt_tokens=prompt_tokens,
@ -320,7 +320,7 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
model_response["choices"][0]["message"] = response_json["message"]
model_response["created"] = int(time.time())
model_response["model"] = "ollama/" + data["model"]
prompt_tokens = response_json["prompt_eval_count"] # type: ignore
prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore
completion_tokens = response_json["eval_count"]
model_response["usage"] = litellm.Usage(
prompt_tokens=prompt_tokens,