From 3df7231fa5aa4cb6637853361925899a6eb3c431 Mon Sep 17 00:00:00 2001 From: frob Date: Fri, 19 Apr 2024 21:38:42 +0200 Subject: [PATCH] Disable special tokens in ollama completion when counting tokens Some(?) models (eg, codegemma) don't return a prompt_eval_count field, so ollama.py tries to compute the value based on encoding of the prompt. Unfortunately FIM symbols used in the prompt (eg, "<|fim_prefix|>") cause the encoder to throw an exception, so we disable special processing. --- litellm/llms/ollama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index a14c3cb5031..96c75232da9 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -228,7 +228,7 @@ def get_ollama_response( model_response["choices"][0]["message"]["content"] = response_json["response"] model_response["created"] = int(time.time()) model_response["model"] = "ollama/" + model - prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt))) # type: ignore + prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(prompt, disallowed_special=()))) # type: ignore completion_tokens = response_json.get("eval_count", len(response_json.get("message",dict()).get("content", ""))) model_response["usage"] = litellm.Usage( prompt_tokens=prompt_tokens, @@ -330,7 +330,7 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): ] model_response["created"] = int(time.time()) model_response["model"] = "ollama/" + data["model"] - prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(data["prompt"]))) # type: ignore + prompt_tokens = response_json.get("prompt_eval_count", len(encoding.encode(data["prompt"], disallowed_special=()))) # type: ignore completion_tokens = response_json.get("eval_count", len(response_json.get("message",dict()).get("content", ""))) model_response["usage"] = litellm.Usage( prompt_tokens=prompt_tokens,