From ab72db3b50bec0b92001801f54736cee25b60cae Mon Sep 17 00:00:00 2001 From: VinciGit00 Date: Mon, 29 Apr 2024 10:26:34 +0200 Subject: [PATCH 1/2] if the conf is not present preset to 8192 the context window --- scrapegraphai/graphs/abstract_graph.py | 19 ++++++++++++------- scrapegraphai/helpers/models_tokens.py | 4 ++++ 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/scrapegraphai/graphs/abstract_graph.py b/scrapegraphai/graphs/abstract_graph.py index 442a809e..aa9bd14c 100644 --- a/scrapegraphai/graphs/abstract_graph.py +++ b/scrapegraphai/graphs/abstract_graph.py @@ -64,13 +64,18 @@ def _create_llm(self, llm_config: dict): 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"]: diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index acb8bea4..dc6d627f 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -31,5 +31,9 @@ "dolphin-mixtral": 32000, "mistral-openorca": 32000, "stablelm-zephyr": 8192 + }, + "gemma": { + "gemma:2b": 8192, + "gemma:7b": 8192 } } From 4f10ac302ec728285610d9c0c0787334a86924a4 Mon Sep 17 00:00:00 2001 From: VinciGit00 Date: Mon, 29 Apr 2024 11:06:30 +0200 Subject: [PATCH 2/2] fixed gemma model --- scrapegraphai/helpers/models_tokens.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py index dc6d627f..6379d763 100644 --- a/scrapegraphai/helpers/models_tokens.py +++ b/scrapegraphai/helpers/models_tokens.py @@ -33,7 +33,6 @@ "stablelm-zephyr": 8192 }, "gemma": { - "gemma:2b": 8192, - "gemma:7b": 8192 + "gemma": 8192, } }