add gpt-4o-mini to llama-index

towardsai · Jul 18, 2024 · 3cf9e02 · 3cf9e02
1 parent 37cbdf5
commit 3cf9e02
Showing 1 changed file with 27 additions and 13 deletions.
diff --git a/scripts/gradio-ui.py b/scripts/gradio-ui.py
@@ -44,19 +44,6 @@
 
 load_dotenv(".env")
 
-GPT4_MODELS.update(
-    {
-        "gpt-4-1106-preview": 128000,
-        "gpt-4-0125-preview": 128000,
-        "gpt-4-turbo-preview": 128000,
-        "gpt-4-turbo-2024-04-09": 128000,
-        "gpt-4-turbo": 128000,
-        "gpt-4o": 128000,
-        "gpt-4o-2024-05-13": 128000,
-        "gpt-4o-mini": 128000,
-        # Add any other models you need
-    }
-)
 
 # from utils import init_mongo_db
 
@@ -113,6 +100,33 @@
     "rag_course",
 ]
 
+
+from llama_index.llms.openai.utils import (
+    ALL_AVAILABLE_MODELS,
+    AZURE_TURBO_MODELS,
+    CHAT_MODELS,
+    GPT3_5_MODELS,
+    GPT3_MODELS,
+    GPT4_MODELS,
+    TURBO_MODELS,
+)
+
+# Add new models to GPT4_MODELS
+new_gpt4_models = {
+    "gpt-4-1106-preview": 128000,
+    "gpt-4-0125-preview": 128000,
+    "gpt-4-turbo-preview": 128000,
+    "gpt-4-turbo-2024-04-09": 128000,
+    "gpt-4-turbo": 128000,
+    "gpt-4o": 128000,
+    "gpt-4o-2024-05-13": 128000,
+    "gpt-4o-mini": 128000,
+}
+GPT4_MODELS.update(new_gpt4_models)
+
+# Update ALL_AVAILABLE_MODELS
+ALL_AVAILABLE_MODELS.update(new_gpt4_models)
+
 # # Initialize MongoDB
 # mongo_db = (
 #     init_mongo_db(uri=MONGODB_URI, db_name="towardsai-buster")