From fa651d4cd9ab8ae9cf58280f1256ceb4171ef088 Mon Sep 17 00:00:00 2001
From: Marco Vinciguerra <mvincig11@gmail.com>
Date: Tue, 6 Aug 2024 15:17:49 +0200
Subject: [PATCH] feat: add grok integration

---
 scrapegraphai/helpers/models_tokens.py | 73 ++++++++++++++------------
 1 file changed, 38 insertions(+), 35 deletions(-)

diff --git a/scrapegraphai/helpers/models_tokens.py b/scrapegraphai/helpers/models_tokens.py
index cb00435d..608c16e4 100644
--- a/scrapegraphai/helpers/models_tokens.py
+++ b/scrapegraphai/helpers/models_tokens.py
@@ -44,41 +44,43 @@ models_tokens = {
         "gemini-1.5-pro-latest": 128000,
         "models/embedding-001": 2048
     },
-    "ollama": { "command-r": 12800, 
-               "codellama": 16000, 
-               "dbrx": 32768, 
-               "deepseek-coder:33b": 16000, 
-               "falcon": 2048, 
-               "llama2": 4096, 
-               "llama3": 8192, 
-               "llama3:70b": 8192,
-               "llama3.1":128000,
-               "llama3.1:70b": 128000,
-               "lama3.1:405b": 128000,
-               "scrapegraph": 8192, 
-               "llava": 4096, 
-               "mixtral:8x22b-instruct": 65536, 
-               "mistral-openorca": 32000, 
-               "nomic-embed-text": 8192, 
-               "nous-hermes2:34b": 4096, 
-               "orca-mini": 2048, 
-               "phi3:3.8b": 12800, 
-               "qwen:0.5b": 32000, 
-               "qwen:1.8b": 32000, 
-               "qwen:4b": 32000, 
-               "qwen:14b": 32000, 
-               "qwen:32b": 32000, 
-               "qwen:72b": 32000, 
-               "qwen:110b": 32000, 
-               "stablelm-zephyr": 8192, 
-               "wizardlm2:8x22b": 65536, 
-               # embedding models
-               "shaw/dmeta-embedding-zh-small-q4": 8192,
-               "shaw/dmeta-embedding-zh-q4": 8192,
-               "chevalblanc/acge_text_embedding": 8192,
-               "martcreation/dmeta-embedding-zh": 8192,
-               "snowflake-arctic-embed": 8192, 
-               "mxbai-embed-large": 512 
+    "ollama": { 
+        "grok-1": 8192,
+        "command-r": 12800, 
+        "codellama": 16000, 
+        "dbrx": 32768, 
+        "deepseek-coder:33b": 16000, 
+        "falcon": 2048, 
+        "llama2": 4096, 
+        "llama3": 8192, 
+        "llama3:70b": 8192,
+        "llama3.1":128000,
+        "llama3.1:70b": 128000,
+        "lama3.1:405b": 128000,
+        "scrapegraph": 8192, 
+        "llava": 4096, 
+        "mixtral:8x22b-instruct": 65536, 
+        "mistral-openorca": 32000, 
+        "nomic-embed-text": 8192, 
+        "nous-hermes2:34b": 4096, 
+        "orca-mini": 2048, 
+        "phi3:3.8b": 12800, 
+        "qwen:0.5b": 32000, 
+        "qwen:1.8b": 32000, 
+        "qwen:4b": 32000, 
+        "qwen:14b": 32000, 
+        "qwen:32b": 32000, 
+        "qwen:72b": 32000, 
+        "qwen:110b": 32000, 
+        "stablelm-zephyr": 8192, 
+        "wizardlm2:8x22b": 65536, 
+        # embedding models
+        "shaw/dmeta-embedding-zh-small-q4": 8192,
+        "shaw/dmeta-embedding-zh-q4": 8192,
+        "chevalblanc/acge_text_embedding": 8192,
+        "martcreation/dmeta-embedding-zh": 8192,
+        "snowflake-arctic-embed": 8192, 
+        "mxbai-embed-large": 512 
     },
     "oneapi": {
         "qwen-turbo": 6000 
@@ -147,6 +149,7 @@ models_tokens = {
         "mistralai/Mistral-7B-Instruct-v0.2": 32000
     },
     "hugging_face": {
+        "xai-org/grok-1": 8192,
         "meta-llama/Meta-Llama-3-8B": 8192,
         "meta-llama/Meta-Llama-3-8B-Instruct": 8192,
         "meta-llama/Meta-Llama-3-70B": 8192,