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

Merge pull request #105 from VinciGit00/fixing_local_models

Browse Source 
if the conf is not present preset to 8192 the context window
This commit is contained in:
Marco Perini 2024-04-29 14:42:23 +02:00 committed by GitHub
commit 8a15f021dc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
scrapegraphai/graphs/abstract_graph.py
View File

@ -64,13 +64,18 @@ class AbstractGraph(ABC):
llm_params["model"] = llm_params["model"].split("/")[-1]
# allow user to set model_tokens in config
if "model_tokens" in llm_params:
self.model_token = llm_params["model_tokens"]
elif llm_params["model"] in models_tokens["ollama"]:
try:
self.model_token = models_tokens["ollama"][llm_params["model"]]
except KeyError:
raise KeyError("Model not supported")
try:
if "model_tokens" in llm_params:
self.model_token = llm_params["model_tokens"]
elif llm_params["model"] in models_tokens["ollama"]:
try:
self.model_token = models_tokens["ollama"][llm_params["model"]]
except KeyError:
raise KeyError("Model not supported")
else:
self.model_token = 8192
except AttributeError:
self.model_token = 8192
return Ollama(llm_params)
elif "hugging_face" in llm_params["model"]:

3
scrapegraphai/helpers/models_tokens.py
View File

@ -31,5 +31,8 @@ models_tokens = {
"dolphin-mixtral": 32000,
"mistral-openorca": 32000,
"stablelm-zephyr": 8192
},
"gemma": {
"gemma": 8192,
}
}