From f79f399ee0d660f162e0cb96d9faba48ecdc88b2 Mon Sep 17 00:00:00 2001 From: Federico Aguzzi <62149513+f-aguzzi@users.noreply.github.com> Date: Mon, 28 Oct 2024 14:51:39 +0100 Subject: [PATCH] fix(AbstractGraph): manually select model tokens closes #768 --- scrapegraphai/graphs/abstract_graph.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index c693400c..bd8244f0 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -152,12 +152,15 @@ def _create_llm(self, llm_config: dict) -> object: raise ValueError(f"""Provider {llm_params['model_provider']} is not supported. If possible, try to use a model instance instead.""") - try: - self.model_token = models_tokens[llm_params["model_provider"]][llm_params["model"]] - except KeyError: - print(f"""Model {llm_params['model_provider']}/{llm_params['model']} not found, - using default token size (8192)""") - self.model_token = 8192 + if "model_tokens" not in llm_params: + try: + self.model_token = models_tokens[llm_params["model_provider"]][llm_params["model"]] + except KeyError: + print(f"""Model {llm_params['model_provider']}/{llm_params['model']} not found, + using default token size (8192)""") + self.model_token = 8192 + else: + self.model_token = llm_params["model_tokens"] try: if llm_params["model_provider"] not in \