Merge pull request #105 from VinciGit00/fixing_local_models

if the conf is not present preset to 8192 the context window

Author: PeriniM

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"]:

scrapegraphai/helpers/models_tokens.py

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