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

fix(AbstractGraph): manually select model tokens

Browse Source 
closes #768
This commit is contained in:
Federico Aguzzi 2024-10-28 14:51:39 +01:00
parent 62369e3e28
commit f79f399ee0
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
scrapegraphai/graphs/abstract_graph.py
View File

@ -152,12 +152,15 @@ class AbstractGraph(ABC):
raise ValueError(f"""Provider {llm_params['model_provider']} is not supported.
If possible, try to use a model instance instead.""")
try:
self.model_token = models_tokens[llm_params["model_provider"]][llm_params["model"]]
except KeyError:
print(f"""Model {llm_params['model_provider']}/{llm_params['model']} not found,
using default token size (8192)""")
self.model_token = 8192
if "model_tokens" not in llm_params:
try:
self.model_token = models_tokens[llm_params["model_provider"]][llm_params["model"]]
except KeyError:
print(f"""Model {llm_params['model_provider']}/{llm_params['model']} not found,
using default token size (8192)""")
self.model_token = 8192
else:
self.model_token = llm_params["model_tokens"]
try:
if llm_params["model_provider"] not in \