O1 azure deployment model class

eureka_ml_insights/configs/model_configs.py

@@ -3,51 +3,74 @@
 You can also add your custom models here by following the same pattern as the existing configs. """
 
 from eureka_ml_insights.models import (
-    ClaudeModels,
-    GeminiModels,
-    LlamaServerlessAzureRestEndpointModels,
-    LLaVA,
-    LLaVAHuggingFaceMM,
-    MistralServerlessAzureRestEndpointModels,
-    OpenAIModelsOAI,
-    RestEndpointModels,
+    AzureOpenAIO1Model,
+    ClaudeModel,
+    DirectOpenAIModel,
+    DirectOpenAIO1Model,
+    GeminiModel,
+    LlamaServerlessAzureRestEndpointModel,
+    LLaVAHuggingFaceModel,
+    LLaVAModel,
+    MistralServerlessAzureRestEndpointModel,
+    RestEndpointModel,
 )
 
 from .config import ModelConfig
 
+# For models that require secret keys, you can store the keys in a json file and provide the path to the file
+# in the secret_key_params dictionary. OR you can provide the key name and key vault URL to fetch the key from Azure Key Vault.
+# You don't need to provide both the key_vault_url and local_keys_path. You can provide one of them based on your setup.
+
 # OpenAI models
 OPENAI_SECRET_KEY_PARAMS = {
     "key_name": "your_openai_secret_key_name",
     "local_keys_path": "keys/keys.json",
     "key_vault_url": None,
 }
 
+OAI_O1_PREVIEW_CONFIG = ModelConfig(
+    DirectOpenAIO1Model,
+    {
+        "model_name": "o1-preview",
+        "secret_key_params": OPENAI_SECRET_KEY_PARAMS,
+    },
+)
+
+OAI_O1_PREVIEW_AUZRE_CONFIG = ModelConfig(
+    AzureOpenAIO1Model,
+    {
+        "model_name": "o1-preview",
+        "url": "your/endpoint/url",
+        "api_version": "2024-08-01-preview",
+    },
+)
+
 OAI_GPT4_1106_PREVIEW_CONFIG = ModelConfig(
-    OpenAIModelsOAI,
+    DirectOpenAIModel,
     {
         "model_name": "gpt-4-1106-preview",
         "secret_key_params": OPENAI_SECRET_KEY_PARAMS,
     },
 )
 
 OAI_GPT4V_1106_VISION_PREVIEW_CONFIG = ModelConfig(
-    OpenAIModelsOAI,
+    DirectOpenAIModel,
     {
         "model_name": "gpt-4-1106-vision-preview",
         "secret_key_params": OPENAI_SECRET_KEY_PARAMS,
     },
 )
 
 OAI_GPT4V_TURBO_2024_04_09_CONFIG = ModelConfig(
-    OpenAIModelsOAI,
+    DirectOpenAIModel,
     {
         "model_name": "gpt-4-turbo-2024-04-09",
         "secret_key_params": OPENAI_SECRET_KEY_PARAMS,
     },
 )
 
 OAI_GPT4O_2024_05_13_CONFIG = ModelConfig(
-    OpenAIModelsOAI,
+    DirectOpenAIModel,
     {
         "model_name": "gpt-4o-2024-05-13",
         "secret_key_params": OPENAI_SECRET_KEY_PARAMS,
@@ -63,15 +86,15 @@
 }
 
 GEMINI_V15_PRO_CONFIG = ModelConfig(
-    GeminiModels,
+    GeminiModel,
     {
         "model_name": "gemini-1.5-pro",
         "secret_key_params": GEMINI_SECRET_KEY_PARAMS,
     },
 )
 
 GEMINI_V1_PRO_CONFIG = ModelConfig(
-    GeminiModels,
+    GeminiModel,
     {
         "model_name": "gemini-1.0-pro",
         "secret_key_params": GEMINI_SECRET_KEY_PARAMS,
@@ -86,15 +109,15 @@
 }
 
 CLAUDE_3_OPUS_CONFIG = ModelConfig(
-    ClaudeModels,
+    ClaudeModel,
     {
         "model_name": "claude-3-opus-20240229",
         "secret_key_params": CLAUDE_SECRET_KEY_PARAMS,
     },
 )
 
 CLAUDE_3_5_SONNET_CONFIG = ModelConfig(
-    ClaudeModels,
+    ClaudeModel,
     {
         "secret_key_params": CLAUDE_SECRET_KEY_PARAMS,
         "model_name": "claude-3-5-sonnet-20240620",
@@ -103,29 +126,29 @@
 
 # LLAVA models
 LLAVAHF_V16_34B_CONFIG = ModelConfig(
-    LLaVAHuggingFaceMM,
+    LLaVAHuggingFaceModel,
     {"model_name": "llava-hf/llava-v1.6-34b-hf", "use_flash_attn": True},
 )
 
 LLAVAHF_V15_7B_CONFIG = ModelConfig(
-    LLaVAHuggingFaceMM,
+    LLaVAHuggingFaceModel,
     {"model_name": "llava-hf/llava-1.5-7b-hf", "use_flash_attn": True},
 )
 
 LLAVA_V16_34B_CONFIG = ModelConfig(
-    LLaVA,
+    LLaVAModel,
     {"model_name": "liuhaotian/llava-v1.6-34b", "use_flash_attn": True},
 )
 
 LLAVA_V15_7B_CONFIG = ModelConfig(
-    LLaVA,
+    LLaVAModel,
     {"model_name": "liuhaotian/llava-v1.5-7b", "use_flash_attn": True},
 )
 
 # Llama models
 
 LLAMA3_1_70B_INSTRUCT_CONFIG = ModelConfig(
-    RestEndpointModels,
+    RestEndpointModel,
     {
         "url": "your/endpoint/url",
         "secret_key_params": {
@@ -138,7 +161,7 @@
 )
 
 LLAMA3_1_405B_INSTRUCT_CONFIG = ModelConfig(
-    LlamaServerlessAzureRestEndpointModels,
+    LlamaServerlessAzureRestEndpointModel,
     {
         "url": "your/endpoint/url",
         "secret_key_params": {
@@ -152,7 +175,7 @@
 
 # Mistral Endpoints
 AIF_NT_MISTRAL_LARGE_2_2407_CONFIG = ModelConfig(
-    MistralServerlessAzureRestEndpointModels,
+    MistralServerlessAzureRestEndpointModel,
     {
         "url": "your/endpoint/url",
         "secret_key_params": {

eureka_ml_insights/models/__init__.py

@@ -1,32 +1,31 @@
 from .models import (
-    ClaudeModels,
-    GeminiModels,
-    HuggingFaceLM,
-    LlamaServerlessAzureRestEndpointModels,
-    LLaVA,
-    LLaVAHuggingFaceMM,
-    MistralServerlessAzureRestEndpointModels,
-    OpenAIModelsMixIn,
-    OpenAIModelsAzure,
-    OpenAIModelsOAI,
-    OpenAIO1Direct,
-    Phi3HF,
-    KeyBasedAuthentication,
-    EndpointModels,
-    RestEndpointModels
+    AzureOpenAIModel,
+    AzureOpenAIO1Model,
+    ClaudeModel,
+    DirectOpenAIModel,
+    DirectOpenAIO1Model,
+    GeminiModel,
+    HuggingFaceModel,
+    LlamaServerlessAzureRestEndpointModel,
+    LLaVAHuggingFaceModel,
+    LLaVAModel,
+    MistralServerlessAzureRestEndpointModel,
+    Phi3HFModel,
+    RestEndpointModel,
 )
 
 __all__ = [
-    OpenAIModelsMixIn,
-    OpenAIO1Direct,
-    HuggingFaceLM,
-    LLaVAHuggingFaceMM,
-    Phi3HF,
-    OpenAIModelsOAI,
-    OpenAIModelsAzure,
-    GeminiModels,
-    ClaudeModels,
-    MistralServerlessAzureRestEndpointModels,
-    LlamaServerlessAzureRestEndpointModels,
-    LLaVA,
+    AzureOpenAIO1Model,
+    DirectOpenAIO1Model,
+    HuggingFaceModel,
+    LLaVAHuggingFaceModel,
+    Phi3HFModel,
+    DirectOpenAIModel,
+    AzureOpenAIModel,
+    GeminiModel,
+    ClaudeModel,
+    MistralServerlessAzureRestEndpointModel,
+    LlamaServerlessAzureRestEndpointModel,
+    LLaVAModel,
+    RestEndpointModel,
 ]