Add MiniMax hosted inference adapter

Author: octo-patch

.github/scripts/spellcheck_conf/wordlist.txt

@@ -1183,6 +1183,8 @@ ShardingStrategy
 hsdp
 prem
 Prem
+Anthropic
+MiniMax
 OpenAI
 Prem
 TCP

requirements.txt

@@ -15,6 +15,7 @@ scipy
 optimum
 matplotlib
 chardet
+anthropic>=0.72.0
 openai
 typing-extensions>=4.8.0
 tabulate

src/README.md

@@ -55,6 +55,31 @@ pip install -U pip setuptools
 pip install -e .[tests,auditnlg,vllm]
 ```
 
+### MiniMax hosted inference
+
+The `MINIMAX` adapter supports `MiniMax-M3` and `MiniMax-M2.7` through both
+compatible API formats. It uses the global OpenAI-compatible endpoint by default.
+
+```python
+import os
+
+from llama_cookbook.inference.llm import MINIMAX
+
+api_key = os.environ["MINIMAX_API_KEY"]
+
+global_openai = MINIMAX("MiniMax-M3", api_key)
+china_openai = MINIMAX(
+    "MiniMax-M2.7", api_key, base_url="https://api.minimaxi.com/v1"
+)
+global_anthropic = MINIMAX("MiniMax-M3", api_key, api_format="anthropic")
+china_anthropic = MINIMAX(
+    "MiniMax-M2.7",
+    api_key,
+    api_format="anthropic",
+    base_url="https://api.minimaxi.com/anthropic",
+)
+```
+
 
 ### Getting the Llama models
 You can find Llama models on Hugging Face hub [here](https://huggingface.co/meta-llama), **where models with `hf` in the name are already converted to Hugging Face checkpoints so no further conversion is needed**. The conversion step below is only for original model weights from Meta that are hosted on Hugging Face model hub as well.
@@ -71,4 +96,4 @@ cd transformers
 pip install protobuf
 python src/transformers/models/llama/convert_llama_weights_to_hf.py \
    --input_dir /path/to/downloaded/llama/weights --model_size 3B --output_dir /output/path
-```
+```

src/llama_cookbook/inference/llm.py

@@ -11,8 +11,9 @@
 
 import time
 from abc import ABC, abstractmethod
-from typing import Callable
+from typing import Callable, Literal
 
+import anthropic
 import openai
 from typing_extensions import override
 
@@ -21,6 +22,11 @@
 TEMPERATURE = 0.1
 TOP_P = 0.9
 
+MINIMAX_DEFAULT_BASE_URLS = {
+    "openai": "https://api.minimax.io/v1",
+    "anthropic": "https://api.minimax.io/anthropic",
+}
+
 LOG: logging.Logger = logging.getLogger(__name__)
 
 
@@ -157,3 +163,66 @@ def valid_models(self) -> list[str]:
             "mistralai/Mistral-7B-Instruct-v0.1",
             "HuggingFaceH4/zephyr-7b-beta",
         ]
+
+
+class MINIMAX(LLM):
+    """Access MiniMax through its OpenAI- or Anthropic-compatible API.
+
+    Global endpoints are used by default. For China endpoints, pass
+    ``https://api.minimaxi.com/v1`` for the OpenAI format or
+    ``https://api.minimaxi.com/anthropic`` for the Anthropic format.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        api_key: str,
+        api_format: Literal["openai", "anthropic"] = "openai",
+        base_url: str | None = None,
+    ) -> None:
+        super().__init__(model, api_key)
+        if api_format not in MINIMAX_DEFAULT_BASE_URLS:
+            raise ValueError(f"Unsupported MiniMax API format: {api_format}")
+
+        self.api_format = api_format
+        resolved_base_url = base_url or MINIMAX_DEFAULT_BASE_URLS[api_format]
+        self.openai_client: openai.OpenAI | None = None
+        self.anthropic_client: anthropic.Anthropic | None = None
+        if api_format == "openai":
+            self.openai_client = openai.OpenAI(
+                base_url=resolved_base_url, api_key=api_key
+            )
+        else:
+            self.anthropic_client = anthropic.Anthropic(
+                base_url=resolved_base_url, api_key=api_key
+            )
+
+    @override
+    def query(self, prompt: str) -> str:
+        level = logging.getLogger().level
+        logging.getLogger().setLevel(logging.WARNING)
+        try:
+            if self.openai_client is not None:
+                response = self.openai_client.chat.completions.create(
+                    model=self.model,
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=MAX_TOKENS,
+                )
+                return response.choices[0].message.content or ""
+
+            if self.anthropic_client is None:
+                raise RuntimeError("MiniMax client is not configured")
+            message = self.anthropic_client.messages.create(
+                model=self.model,
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=MAX_TOKENS,
+            )
+            return "".join(
+                block.text for block in message.content if block.type == "text"
+            )
+        finally:
+            logging.getLogger().setLevel(level)
+
+    @override
+    def valid_models(self) -> list[str]:
+        return ["MiniMax-M3", "MiniMax-M2.7"]

src/tests/test_llm.py

@@ -0,0 +1,119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import unittest
+from unittest.mock import patch
+
+import anthropic
+import httpx
+import openai
+
+from llama_cookbook.inference.llm import MINIMAX
+
+OPENAI_CLIENT = openai.OpenAI
+ANTHROPIC_CLIENT = anthropic.Anthropic
+
+
+class MiniMaxTest(unittest.TestCase):
+    def test_supported_models_and_endpoint_paths(self) -> None:
+        cases = [
+            (
+                "openai",
+                "MiniMax-M3",
+                None,
+                "https://api.minimax.io/v1/chat/completions",
+            ),
+            (
+                "openai",
+                "MiniMax-M2.7",
+                "https://api.minimaxi.com/v1",
+                "https://api.minimaxi.com/v1/chat/completions",
+            ),
+            (
+                "anthropic",
+                "MiniMax-M3",
+                None,
+                "https://api.minimax.io/anthropic/v1/messages",
+            ),
+            (
+                "anthropic",
+                "MiniMax-M2.7",
+                "https://api.minimaxi.com/anthropic",
+                "https://api.minimaxi.com/anthropic/v1/messages",
+            ),
+        ]
+
+        for api_format, model, base_url, expected_url in cases:
+            with self.subTest(api_format=api_format, base_url=base_url):
+                requests = []
+
+                def handler(request: httpx.Request) -> httpx.Response:
+                    requests.append(request)
+                    if api_format == "openai":
+                        body = {
+                            "id": "chatcmpl-test",
+                            "object": "chat.completion",
+                            "created": 0,
+                            "model": model,
+                            "choices": [
+                                {
+                                    "index": 0,
+                                    "message": {
+                                        "role": "assistant",
+                                        "content": "Test response",
+                                    },
+                                    "finish_reason": "stop",
+                                }
+                            ],
+                        }
+                    else:
+                        body = {
+                            "id": "msg_test",
+                            "type": "message",
+                            "role": "assistant",
+                            "content": [{"type": "text", "text": "Test response"}],
+                            "model": model,
+                            "stop_reason": "end_turn",
+                            "stop_sequence": None,
+                            "usage": {"input_tokens": 1, "output_tokens": 1},
+                        }
+                    return httpx.Response(200, json=body, request=request)
+
+                http_client = httpx.Client(transport=httpx.MockTransport(handler))
+                kwargs = {"api_format": api_format}
+                if base_url is not None:
+                    kwargs["base_url"] = base_url
+
+                if api_format == "openai":
+                    client = OPENAI_CLIENT(
+                        api_key="test-key",
+                        base_url=base_url or "https://api.minimax.io/v1",
+                        http_client=http_client,
+                    )
+                    patch_target = "llama_cookbook.inference.llm.openai.OpenAI"
+                else:
+                    client = ANTHROPIC_CLIENT(
+                        api_key="test-key",
+                        base_url=base_url or "https://api.minimax.io/anthropic",
+                        http_client=http_client,
+                    )
+                    patch_target = "llama_cookbook.inference.llm.anthropic.Anthropic"
+
+                with patch(patch_target, return_value=client):
+                    provider = MINIMAX(model, "test-key", **kwargs)
+                    self.assertEqual(provider.query("Test prompt"), "Test response")
+                    self.assertEqual(
+                        provider.valid_models(), ["MiniMax-M3", "MiniMax-M2.7"]
+                    )
+
+                self.assertEqual(len(requests), 1)
+                self.assertEqual(str(requests[0].url), expected_url)
+                self.assertEqual(json.loads(requests[0].content)["model"], model)
+                client.close()
+
+
+if __name__ == "__main__":
+    unittest.main()