Skip to content

Commit ffd6015

Browse files
authored
Update abstract_graph.py
Adapt to some ollama modles with "nam/modelname" like ( 893379029/piccolo-large-zh-v2)
1 parent eb41f0d commit ffd6015

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

scrapegraphai/graphs/abstract_graph.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ def _create_llm(self, llm_config: dict, chat=False) -> object:
157157
raise KeyError("Model not supported") from exc
158158
return Anthropic(llm_params)
159159
elif "ollama" in llm_params["model"]:
160-
llm_params["model"] = llm_params["model"].split("/")[-1]
160+
llm_params["model"] = llm_params["model"].split("ollama/")[-1]
161161

162162
# allow user to set model_tokens in config
163163
try:
@@ -273,7 +273,7 @@ def _create_embedder(self, embedder_config: dict) -> object:
273273
elif "azure" in embedder_config["model"]:
274274
return AzureOpenAIEmbeddings()
275275
elif "ollama" in embedder_config["model"]:
276-
embedder_config["model"] = embedder_config["model"].split("/")[-1]
276+
embedder_config["model"] = embedder_config["model"].split("ollama/")[-1]
277277
try:
278278
models_tokens["ollama"][embedder_config["model"]]
279279
except KeyError as exc:

0 commit comments

Comments
 (0)