fix: abstract graph local model

This commit is contained in:
Marco Vinciguerra 2024-08-26 18:00:17 +02:00
parent c1ce9c69d4
commit 04128e7e9f
2 changed files with 3 additions and 2 deletions

View File

@ -170,8 +170,8 @@ class AbstractGraph(ABC):
elif "ollama" in llm_params["model"]:
model_name = llm_params["model"].split("ollama/")[-1]
token_key = model_name if "model_tokens" not in llm_params else None
explicit_model_tokens = 8192 if "model_tokens" not in llm_params else llm_params["model_tokens"]
return handle_model(model_name, "ollama", token_key, explicit_model_tokens)
model_tokens = 8192 if "model_tokens" not in llm_params else llm_params["model_tokens"]
return handle_model(model_name, "ollama", token_key, model_tokens)
elif "claude-3-" in llm_params["model"]:
return handle_model(llm_params["model"], "anthropic", "claude3")

View File

@ -74,6 +74,7 @@ models_tokens = {
"llama3.1:70b": 128000,
"lama3.1:405b": 128000,
"scrapegraph": 8192,
"mistral": 8192,
"llava": 4096,
"mixtral:8x22b-instruct": 65536,
"mistral-openorca": 32000,