Add Clarifai as a provider in MLflow Deployments Server

docs/source/llms/deployments/index.rst

@@ -40,8 +40,8 @@ organizations that use LLMs on a regular basis.
 Tutorials and Guides
 ====================
 
-If you're interested in diving right in to a step by step guide that will get you up and running with the MLflow Deployments Server 
-as fast as possible, the guides below will be your best first stop. 
+If you're interested in diving right in to a step by step guide that will get you up and running with the MLflow Deployments Server
+as fast as possible, the guides below will be your best first stop.
 
 .. raw:: html
 
@@ -256,6 +256,9 @@ below can be used as a helpful guide when configuring a given endpoint for any n
 | AWS Bedrock              | - Amazon Titan           | N/A                      | N/A                      |
 |                          | - Third-party providers  |                          |                          |
 +--------------------------+--------------------------+--------------------------+--------------------------+
+| Clarifai                 | - Clarifai models        | N/A                      | - Clarifai models        |
+|                          | - Third-party providers  |                          | - Third-party providers  |
++--------------------------+--------------------------+--------------------------+--------------------------+
 
 
 † Llama 2 is licensed under the `LLAMA 2 Community License <https://ai.meta.com/llama/license/>`_, Copyright © Meta Platforms, Inc. All Rights Reserved.
@@ -297,6 +300,7 @@ As of now, the MLflow Deployments Server supports the following providers:
 * **huggingface text generation inference**: This is used for models deployed using `Huggingface Text Generation Inference <https://huggingface.co/docs/text-generation-inference/index>`_.
 * **ai21labs**: This is used for models offered by `AI21 Labs <https://studio.ai21.com/foundation-models>`_.
 * **bedrock**: This is used for models offered by `AWS Bedrock <https://aws.amazon.com/bedrock/>`_.
+* **clarifai**: This is used for models offered by `Clarifai <https://clarifai.com/explore/models/>`_.
 
 More providers are being added continually. Check the latest version of the MLflow Deployments Server Docs for the
 most up-to-date list of supported providers.
@@ -494,6 +498,7 @@ Each endpoint has the following configuration parameters:
     - "huggingface-text-generation-inference"
     - "ai21labs"
     - "bedrock"
+    - "clarifai"
 
   - **name**: This is an optional field to specify the name of the model.
   - **config**: This contains provider-specific configuration details.
@@ -580,6 +585,22 @@ Anthropic
 | **anthropic_api_key**   | Yes      | N/A                      | This is the API key for the Anthropic service.        |
 +-------------------------+----------+--------------------------+-------------------------------------------------------+
 
+Clarifai
+++++++++
+
++-------------------------+----------+--------------------------+-------------------------------------------------------+
+| Configuration Parameter | Required | Default                  | Description                                           |
++=========================+==========+==========================+=======================================================+
+| **clarifai_pat**        | Yes      | N/A                      | This is the personal access token for the Clarifai API|
++-------------------------+----------+--------------------------+-------------------------------------------------------+
+| **user_id**             | Yes      | N/A                      | user ID for the respective Clarifai model             |
++-------------------------+----------+--------------------------+-------------------------------------------------------+
+| **app_id**              | Yes      | N/A                      | app ID for the respective Clarifai model              |
++-------------------------+----------+--------------------------+-------------------------------------------------------+
+| **model_version_id**    | No       | N/A                      | model version ID for the respective model.            |
+|                         |          |                          | If not provided selects the latest version by default |
++-------------------------+----------+--------------------------+-------------------------------------------------------+
+
 AWS Bedrock
 +++++++++++
 

examples/deployments/deployments_server/clarifai/README.md

@@ -0,0 +1,57 @@
+## Example endpoint configuration for Clarifai
+
+\
+[Clarifai](https://clarifai.com/) enables easy and efficient deployment of cutting-edge large language models (LLMs) for text processing. Developers can leverage Clarifai's platform to serve up predictions from state-of-the-art open source LLMs including options like [Llama2](https://clarifai.com/meta/Llama-2/models/llama2-7b-chat), [Falcon](https://clarifai.com/tiiuae/falcon/models/falcon-40b-instruct), [Mistral](https://clarifai.com/mistralai/completion/models/mistral-7B-Instruct), [WizardLM](https://clarifai.com/explore/models?searchQuery=wizard&page=1&perPage=24), [Zephyr](https://clarifai.com/huggingface-research/zephyr/models/zephyr-7B-alpha) and many more for text generation.
+
+Explore more Clarifai hosted LLMs [here](https://clarifai.com/explore/models).
+
+### Getting Started
+
+#### Setting a Clarifai PAT
+
+This example requires a [Clarifai PAT](https://docs.clarifai.com/clarifai-basics/authentication/personal-access-tokens/):
+
+```sh
+export CLARIFAI_PAT=...
+```
+
+### Update the config.yaml
+
+To configure access to Clarifai LLMs, update the MLflow Desployments server `config.yaml` file. The example below shows how to specify routes for both text completions and text embeddings using Clarifai models:
+
+```yaml
+endpoints:
+  - name: completions
+    endpoint_type: llm/v1/completions
+    model:
+      provider: clarifai
+      name: mistral-7B-Instruct
+      config:
+        CLARIFAI_PAT: $CLARIFAI_PAT
+        user_id: mistralai
+        app_id: completion
+
+  - name: embeddings
+    endpoint_type: llm/v1/embeddings
+    model:
+      provider: clarifai
+      name: multimodal-clip-embed
+      config:
+        CLARIFAI_PAT: $CLARIFAI_PAT
+        user_id: clarifai
+        app_id: main
+```
+
+This configuration file specifies two routes: 'completions' and 'embeddings', both using models 'mistral-7B-Instruct' and 'multimodal-clip-embed' hosted in Clarifai Platform, respectively.
+
+### Starting the MLflow Deployments server
+
+After the configuration file is created, you can start the gateway by running the following command:
+
+```
+mlflow deployments start-server --config-path examples/deployments/deployments_server/clarifai/config.yaml --port 7000
+```
+
+### Querying the MLflow Deployments server
+
+See the [example script](example.py) within this directory to see how to query the `mistral-7B-Instruct` and `multimodal-clip-embed` model that is served.

examples/deployments/deployments_server/clarifai/config.yaml

@@ -0,0 +1,20 @@
+endpoints:
+  - name: completions
+    endpoint_type: llm/v1/completions
+    model:
+      provider: clarifai
+      name: mistral-7B-Instruct
+      config:
+        CLARIFAI_PAT: $CLARIFAI_PAT
+        user_id: mistralai
+        app_id: completion
+
+  - name: embeddings
+    endpoint_type: llm/v1/embeddings
+    model:
+      provider: clarifai
+      name: multimodal-clip-embed
+      config:
+        CLARIFAI_PAT: $CLARIFAI_PAT
+        user_id: clarifai
+        app_id: main

examples/deployments/deployments_server/clarifai/example.py

@@ -0,0 +1,30 @@
+from mlflow.deployments import get_deploy_client
+
+
+def main():
+    # Set the URI for the MLflow AI Gateway
+    client = get_deploy_client("http://localhost:7000")
+
+    print(f"Clarifai endpoints: {client.list_endpoints()}\n")
+    print(f"Clarifai completions endpoint info: {client.get_endpoint(endpoint='completions')}\n")
+
+    # Completions request
+    response_completions = client.predict(
+        endpoint="completions",
+        inputs={
+            "prompt": "<s><INST>What are some economic impacts that can occur due to seasonal changes in different industries?</INST>",
+            "temperature": 0.7,
+        },
+    )
+    print(f"Clarifai response for completions: {response_completions}")
+
+    # Embeddings request
+    response_embeddings = client.predict(
+        endpoint="embeddings",
+        inputs={"input": ["Do you carry the Storm Trooper costume in size 2T?"]},
+    )
+    print(f"Clarifai response for embeddings: {response_embeddings}")
+
+
+if __name__ == "__main__":
+    main()

mlflow/gateway/config.py

@@ -46,6 +46,7 @@ class Provider(str, Enum):
     # Note: The following providers are only supported on Databricks
     DATABRICKS_MODEL_SERVING = "databricks-model-serving"
     DATABRICKS = "databricks"
+    CLARIFAI = "clarifai"
 
     @classmethod
     def values(cls):
@@ -215,6 +216,18 @@ class AWSBedrockConfig(ConfigModel):
     aws_config: Union[AWSRole, AWSIdAndKey, AWSBaseConfig]
 
 
+class ClarifaiConfig(ConfigModel):
+    CLARIFAI_PAT: str
+    user_id: str
+    app_id: str
+    model_version_id: Optional[str] = None
+
+    # pylint: disable=no-self-argument
+    @validator("CLARIFAI_PAT", pre=True)
+    def validate_CLARIFAI_PAT(cls, value):
+        return _resolve_api_key_from_input(value)
+
+
 config_types = {
     Provider.COHERE: CohereConfig,
     Provider.OPENAI: OpenAIConfig,
@@ -225,6 +238,7 @@ class AWSBedrockConfig(ConfigModel):
     Provider.MLFLOW_MODEL_SERVING: MlflowModelServingConfig,
     Provider.PALM: PaLMConfig,
     Provider.HUGGINGFACE_TEXT_GENERATION_INFERENCE: HuggingFaceTextGenerationInferenceConfig,
+    Provider.CLARIFAI: ClarifaiConfig,
 }
 
 
@@ -284,6 +298,7 @@ class Model(ConfigModel):
             MlflowModelServingConfig,
             HuggingFaceTextGenerationInferenceConfig,
             PaLMConfig,
+            ClarifaiConfig,
         ]
     ] = None
 

mlflow/gateway/providers/__init__.py

@@ -9,6 +9,7 @@ def get_provider(provider: Provider) -> Type[BaseProvider]:
     from mlflow.gateway.providers.ai21labs import AI21LabsProvider
     from mlflow.gateway.providers.anthropic import AnthropicProvider
     from mlflow.gateway.providers.bedrock import AWSBedrockProvider
+    from mlflow.gateway.providers.clarifai import ClarifaiProvider
     from mlflow.gateway.providers.cohere import CohereProvider
     from mlflow.gateway.providers.huggingface import HFTextGenerationInferenceServerProvider
     from mlflow.gateway.providers.mlflow import MlflowModelServingProvider
@@ -26,6 +27,7 @@ def get_provider(provider: Provider) -> Type[BaseProvider]:
         Provider.MLFLOW_MODEL_SERVING: MlflowModelServingProvider,
         Provider.HUGGINGFACE_TEXT_GENERATION_INFERENCE: HFTextGenerationInferenceServerProvider,
         Provider.BEDROCK: AWSBedrockProvider,
+        Provider.CLARIFAI: ClarifaiProvider,
     }
     if prov := provider_to_class.get(provider):
         return prov