From 03ffebc52de3fc6f80a968880e8ade3e3cdf95ec Mon Sep 17 00:00:00 2001 From: Marco Vinciguerra Date: Mon, 17 Jun 2024 14:17:44 +0200 Subject: [PATCH] fix: add chinese embedding model --- scrapegraphai/graphs/abstract_graph.py | 2 +- scrapegraphai/helpers/models_tokens.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index b5e15e8f..20e61e2d 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -320,7 +320,7 @@ class AbstractGraph(ABC): elif "azure" in embedder_params["model"]: return AzureOpenAIEmbeddings() elif "ollama" in embedder_params["model"]: - embedder_params["model"] = embedder_params["model"].split("ollama/")[-1] + embedder_params["model"] = "/".join(embedder_params["model"].split("/")[1:]) try: models_tokens["ollama"][embedder_params["model"]] except KeyError as exc: diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index 0ec94795..c9d61a98 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -59,7 +59,8 @@ models_tokens = { "qwen:110b": 32000, "stablelm-zephyr": 8192, "wizardlm2:8x22b": 65536, - # embedding models + # embedding models + "shaw/dmeta-embedding-zh": 8192, "snowflake-arctic-embed": 8192, "mxbai-embed-large": 512 },