CodeCow/Scrapegraph-ai
1
0
Fork 0
mirror of https://github.com/VinciGit00/Scrapegraph-ai.git synced 2026-07-12 21:01:56 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: shallow copy config of create_embedder

Browse Source
This commit is contained in:
liaoliaojun 2024-06-15 17:52:55 +08:00
parent c44b701d2d
commit 62b372b675
1 changed files with 20 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
scrapegraphai/graphs/abstract_graph.py
View File

@ -333,40 +333,41 @@ class AbstractGraph(ABC):
Raises: Raises:
KeyError: If the model is not supported. KeyError: If the model is not supported.
""" """
embedder_params = {**embedder_config}
if "model_instance" in 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 # Instantiate the embedding model based on the model name
if "openai" in embedder_config["model"]: if "openai" in embedder_params["model"]:
return OpenAIEmbeddings(api_key=embedder_config["api_key"]) return OpenAIEmbeddings(api_key=embedder_params["api_key"])
elif "azure" in embedder_config["model"]: elif "azure" in embedder_params["model"]:
return AzureOpenAIEmbeddings() return AzureOpenAIEmbeddings()
elif "ollama" in embedder_config["model"]: elif "ollama" in embedder_params["model"]:
embedder_config["model"] = embedder_config["model"].split("ollama/")[-1] embedder_params["model"] = embedder_params["model"].split("ollama/")[-1]
try: try:
models_tokens["ollama"][embedder_config["model"]] models_tokens["ollama"][embedder_params["model"]]
except KeyError as exc: except KeyError as exc:
raise KeyError("Model not supported") from exc raise KeyError("Model not supported") from exc
return OllamaEmbeddings(**embedder_config) return OllamaEmbeddings(**embedder_params)
elif "hugging_face" in embedder_config["model"]: elif "hugging_face" in embedder_params["model"]:
try: try:
models_tokens["hugging_face"][embedder_config["model"]] models_tokens["hugging_face"][embedder_params["model"]]
except KeyError as exc: except KeyError as exc:
raise KeyError("Model not supported") from exc raise KeyError("Model not supported") from exc
return HuggingFaceHubEmbeddings(model=embedder_config["model"]) return HuggingFaceHubEmbeddings(model=embedder_params["model"])
elif "gemini" in embedder_config["model"]: elif "gemini" in embedder_params["model"]:
try: try:
models_tokens["gemini"][embedder_config["model"]] models_tokens["gemini"][embedder_params["model"]]
except KeyError as exc: except KeyError as exc:
raise KeyError("Model not supported") from exc raise KeyError("Model not supported") from exc
return GoogleGenerativeAIEmbeddings(model=embedder_config["model"]) return GoogleGenerativeAIEmbeddings(model=embedder_params["model"])
elif "bedrock" in embedder_config["model"]: elif "bedrock" in embedder_params["model"]:
embedder_config["model"] = embedder_config["model"].split("/")[-1] embedder_params["model"] = embedder_params["model"].split("/")[-1]
client = embedder_config.get("client", None) client = embedder_params.get("client", None)
try: try:
models_tokens["bedrock"][embedder_config["model"]] models_tokens["bedrock"][embedder_params["model"]]
except KeyError as exc: except KeyError as exc:
raise KeyError("Model not supported") from 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: else:
raise ValueError("Model provided by the configuration not supported") raise ValueError("Model provided by the configuration not supported")