mistral-inference

repository·main·Indexed 27 days ago

https://github.com/mistralai/mistral-inference

A library providing minimal code for high-performance local inference of Mistral models, including Mistral 7B, Mixtral (MoE), Codestral, and others. It includes tools for interactive chat, model demos, and Python implementations for instruction following, multimodal tasks, function calling, and Fill-in-the-middle (FIM). The package supports KV caching via BufferCache and CacheView, and provides configuration classes like TransformerArgs, MambaArgs, and VisionEncoderArgs.

Tokens
8.1K
Snippets
16
Records
39
Agent score
92%

What's inside mistral-inference

  1. Install mistral-inference from local source

    main

    Clone the repository and use poetry to install the dependencies locally. This also requires a GPU for the xformers dependency.

    cd $HOME && git clone https://github.com/mistralai/mistral-inference
    cd $HOME/mistral-inference && poetry install .
  2. Build a vLLM deployment image

    main

    The deploy folder contains instructions to build a Docker image for serving Mistral AI models via vLLM. This image uses the transformers library instead of the reference implementation. Build the image using the following command:

    docker build deploy --build-arg MAX_JOBS=8
  3. Download Mistral models via direct links

    main

    Mistral models can be downloaded directly from mistralcdn.com. After downloading, extract the .tar files into a dedicated model directory.

    Available Models (Direct):

    • 7B Instruct: https://models.mistralcdn.com/mistral-7b-v0-3/mistral-7B-Instruct-v0.3.tar
    • 8x7B Instruct: https://models.mistralcdn.com/mixtral-8x7b-v0-1/Mixtral-8x7B-v0.1-Instruct.tar
    • 8x22 Instruct: https://models.mistralcdn.com/mixtral-8x22b-v0-3/mixtral-8x22B-Instruct-v0.3.tar
    • 7B Base: https://models.mistralcdn.com/mistral-7b-v0-3/mistral-7B-v0.3.tar
    • 8x22B: https://models.mistralcdn.com/mixtral-8x22b-v0-3/mixtral-8x22B-v0.3.tar
    • Codestral 22B: https://models.mistralcdn.com/codestral-22b-v0-1/codestral-22B-v0.1.tar
    • Mathstral 7B: https://models.mistralcdn.com/mathstral-7b-v0-1/mathstral-7B-v0.1.tar
    • Codestral-Mamba 7B: https://models.mistralcdn.com/codestral-mamba-7b-v0-1/codestral-mamba-7B-v0.1.tar
    • Nemo Base: https://models.mistralcdn.com/mistral-nemo-2407/mistral-nemo-base-2407.tar
    • Nemo Instruct: https://models.mistralcdn.com/mistral-nemo-2407/mistral-nemo-instruct-2407.tar
    • Mistral Large 2: https://models.mistralcdn.com/mistral-large-2407/mistral-large-instruct-2407.tar
    export MISTRAL_MODEL=$HOME/mistral_models
    mkdir -p $MISTRAL_MODEL
    
    export 12B_DIR=$MISTRAL_MODEL/12B_Nemo
    wget https://models.mistralcdn.com/mistral-nemo-2407/mistral-nemo-instruct-2407.tar
    mkdir -p $12B_DIR
    tar -xf mistral-nemo-instruct-2407.tar -C $12B_DIR
  4. Train a classifier on Mistral features

    main

    A high-performance way to classify text is to:

    1. Use Mistral to generate embeddings for your dataset using forward_partial.
    2. Average the token embeddings to get a single vector per sample.
    3. Train a standard machine learning classifier (like sklearn.linear_model.LogisticRegression) on these frozen features.
    4. Normalize features using sklearn.preprocessing.StandardScaler for better stability.
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import StandardScaler
    
    # 1. Normalize
    scaler = StandardScaler()
    train_x = scaler.fit_transform(train_x)
    
    # 2. Train
    clf = LogisticRegression(random_state=0, C=1.0, max_iter=500).fit(train_x, train_y)
  5. Manage BufferCache lifecycle and sequence lengths

    main

    To use a BufferCache, you must manage the sequence lengths tracked within it:

    1. Initialize sequence lengths: Call init_kvseqlens(batch_size) to set up the internal tracking tensor.
    2. Update sequence lengths: As new tokens are processed, call update_seqlens(seqlens) where seqlens is a List[int] representing the number of new tokens added to each sequence in the batch.
    3. Reset: Call reset() to clear the tracked sequence lengths.
    4. Device/Dtype: Use .to(device, dtype) to move the cache to a specific device or change its precision.
  6. Perform zero-shot classification with prompting

    main
    Zero-shot classification can be achieved by prompting the model with a specific template (e.g., Symptoms: {symptom}\nDisease:) and evaluating the log-probabilities of each possible label appearing immediately after the prompt. This method does not require training but is significantly slower and typically less accurate than training a linear classifier on top of frozen features.
  7. Download Mistral models from Hugging Face Hub

    main

    You can download Mistral models from the Hugging Face Hub using huggingface_hub.snapshot_download. This is recommended for models like Pixtral and Mistral Small 3.1.

    Hugging Face IDs:

    • Pixtral Large Instruct: mistralai/Pixtral-Large-Instruct-2411
    • Pixtral 12B Base: mistralai/Pixtral-12B-Base-2409
    • Pixtral 12B: mistralai/Pixtral-12B-2409
    • Mistral Small 3.1 24B Base: mistralai/Mistral-Small-3.1-24B-Base-2503
    • Mistral Small 3.1 24B Instruct: mistralai/Mistral-Small-3.1-24B-Instruct-2503
    from pathlib import Path
    from huggingface_hub import snapshot_download
    
    mistral_models_path = Path.home().joinpath("mistral_models")
    
    model_path = mistral_models_path / "mistral-small-3.1-instruct"
    model_path.mkdir(parents=True, exist_ok=True)
    
    repo_id = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
    
    snapshot_download(
        repo_id=repo_id,
        allow_patterns=["params.json", "consolidated.safetensors", "tekken.json"],
        local_dir=model_path,
    )
  8. Implement Function Calling in Python

    main

    To use function calling, define your tools using Tool and Function objects from mistral_common.protocol.instruct.tool_calls. Include the function name, description, and a JSON schema for parameters. Pass these tools into the ChatCompletionRequest before encoding.

    from mistral_common.protocol.instruct.tool_calls import Function, Tool
    from mistral_common.protocol.instruct.request import ChatCompletionRequest
    from mistral_common.protocol.instruct.messages import UserMessage
    
    completion_request = ChatCompletionRequest(
        tools=[
            Tool(
                function=Function(
                    name="get_current_weather",
                    description="Get the current weather",
                    parameters={
                        "type": "object",
                        "properties": {
                            "location": {
                                "type": "string",
                                "description": "The city and state, e.g. San Francisco, CA",
                            },
                            "format": {
                                "type": "string",
                                "enum": ["celsius", "fahrenheit"],
                                "description": "The temperature unit to use. Infer this from the users location.",
                            },
                        },
                        "required": ["location", "format"],
                    },
                )
            )
        ],
        messages=[
            UserMessage(content="What's the weather like today in Paris?"),
        ],
    )
    
    tokens = tokenizer.encode_chat_completion(completion_request).tokens
    
    out_tokens, _ = generate([tokens], model, max_tokens=64, temperature=0.0, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
    result = tokenizer.instruct_tokenizer.tokenizer.decode(out_tokens[0])
    
    print(result)