From 62b372b675a45ca4d031f337b6f8728151689442 Mon Sep 17 00:00:00 2001 From: liaoliaojun Date: Sat, 15 Jun 2024 17:52:55 +0800 Subject: [PATCH] fix: shallow copy config of create_embedder --- scrapegraphai/graphs/abstract_graph.py | 39 +++++++++++++------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index b5f3a681..49a6cb5f 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -333,40 +333,41 @@ class AbstractGraph(ABC): Raises: KeyError: If the model is not supported. """ + embedder_params = {**embedder_config} if "model_instance" in embedder_config: - return embedder_config["model_instance"] + return embedder_params["model_instance"] # Instantiate the embedding model based on the model name - if "openai" in embedder_config["model"]: - return OpenAIEmbeddings(api_key=embedder_config["api_key"]) - elif "azure" in embedder_config["model"]: + if "openai" in embedder_params["model"]: + return OpenAIEmbeddings(api_key=embedder_params["api_key"]) + elif "azure" in embedder_params["model"]: return AzureOpenAIEmbeddings() - elif "ollama" in embedder_config["model"]: - embedder_config["model"] = embedder_config["model"].split("ollama/")[-1] + elif "ollama" in embedder_params["model"]: + embedder_params["model"] = embedder_params["model"].split("ollama/")[-1] try: - models_tokens["ollama"][embedder_config["model"]] + models_tokens["ollama"][embedder_params["model"]] except KeyError as exc: raise KeyError("Model not supported") from exc - return OllamaEmbeddings(**embedder_config) - elif "hugging_face" in embedder_config["model"]: + return OllamaEmbeddings(**embedder_params) + elif "hugging_face" in embedder_params["model"]: try: - models_tokens["hugging_face"][embedder_config["model"]] + models_tokens["hugging_face"][embedder_params["model"]] except KeyError as exc: raise KeyError("Model not supported") from exc - return HuggingFaceHubEmbeddings(model=embedder_config["model"]) - elif "gemini" in embedder_config["model"]: + return HuggingFaceHubEmbeddings(model=embedder_params["model"]) + elif "gemini" in embedder_params["model"]: try: - models_tokens["gemini"][embedder_config["model"]] + models_tokens["gemini"][embedder_params["model"]] except KeyError as exc: raise KeyError("Model not supported") from exc - return GoogleGenerativeAIEmbeddings(model=embedder_config["model"]) - elif "bedrock" in embedder_config["model"]: - embedder_config["model"] = embedder_config["model"].split("/")[-1] - client = embedder_config.get("client", None) + return GoogleGenerativeAIEmbeddings(model=embedder_params["model"]) + elif "bedrock" in embedder_params["model"]: + embedder_params["model"] = embedder_params["model"].split("/")[-1] + client = embedder_params.get("client", None) try: - models_tokens["bedrock"][embedder_config["model"]] + models_tokens["bedrock"][embedder_params["model"]] except KeyError as exc: raise KeyError("Model not supported") from exc - return BedrockEmbeddings(client=client, model_id=embedder_config["model"]) + return BedrockEmbeddings(client=client, model_id=embedder_params["model"]) else: raise ValueError("Model provided by the configuration not supported")