From 04128e7e9f585aaf774fabf646c4d9d3b96b8333 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Mon, 26 Aug 2024 18:00:17 +0200 Subject: [PATCH] fix: abstract graph local model --- scrapegraphai/graphs/abstract_graph.py | 4 ++-- scrapegraphai/helpers/models_tokens.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index 71773ff6..2dcc54f9 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -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") diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index 1e3e1910..7677a901 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -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,