fix: gemini errror

VinciGit00 · May 2, 2024 · 2ea54ea · 2ea54ea
1 parent 35ae76f
commit 2ea54ea
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 13 deletions.
diff --git a/examples/gemini/csv_scraper_gemini.py b/examples/gemini/csv_scraper_gemini.py
@@ -19,20 +19,13 @@
 # ************************************************
 # Define the configuration for the graph
 # ************************************************
+gemini_key = os.getenv("GOOGLE_APIKEY")
 
 graph_config = {
     "llm": {
-        "model": "ollama/mistral",
-        "temperature": 0,
-        "format": "json",  # Ollama needs the format to be specified explicitly
-        # "model_tokens": 2000, # set context length arbitrarily
-        "base_url": "http://localhost:11434",
+        "api_key": gemini_key,
+        "model": "gemini-pro",
     },
-    "embeddings": {
-        "model": "ollama/nomic-embed-text",
-        "temperature": 0,
-        "base_url": "http://localhost:11434",
-    }
 }
 
 # ************************************************

diff --git a/examples/gemini/smart_scraper_gemini.py b/examples/gemini/smart_scraper_gemini.py
@@ -18,7 +18,7 @@
 graph_config = {
     "llm": {
         "api_key": gemini_key,
-        "model": "gpt-3.5-turbo",
+        "model": "gemini-pro",
     },
 }
 

diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py
@@ -20,7 +20,9 @@
         "gpt-4-32k-0613": 32768,
     },
     "azure": {
-        "gpt-3.5-turbo": 4096
+        "gpt-3.5-turbo": 4096,
+        "gpt-4": 8192,
+        "gpt-4-32k": 32768
     },
     "gemini": {
         "gemini-pro": 128000,
@@ -48,4 +50,3 @@
         "claude3": 200000
     }
 }
-