From df70b4f75d60cda2dddff3c707dc0edb34b80ee6 Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Tue, 27 Aug 2024 18:08:53 +0200 Subject: [PATCH] Update abstract_graph.py --- scrapegraphai/graphs/abstract_graph.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index 58eb30f4..067af7d4 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -125,19 +125,19 @@ class AbstractGraph(ABC): self.model_token = llm_params["model_tokens"] except KeyError as exc: raise KeyError("model_tokens not specified") from exc - return llm_params["model_instance"] + return llm_params["model_instance"] known_providers = {"openai", "azure_openai", "google_genai", "google_vertexai", "ollama", "oneapi", "nvidia", "groq", "anthropic" "bedrock", "mistralai", "hugging_face", "deepseek", "ernie", "fireworks"} - + split_model_provider = llm_params["model"].split("/") llm_params["model_provider"] = split_model_provider[0] llm_params["model"] = split_model_provider[1:] if llm_params["model_provider"] not in known_providers: 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"]].get(llm_params["model"][0]) except KeyError: