Text Embeddings Inference (TEI)

repository·main·Indexed 26 days ago

https://github.com/huggingface/text-embeddings-inference

A high-performance toolkit for deploying and serving open-source text embeddings and sequence classification models. TEI supports various hardware architectures including CPU, CUDA GPUs (Turing, Ampere, Ada Lovelace, Hopper, Blackwell), Apple Silicon, and HPU. It provides endpoints for text embeddings (/embed), re-ranking (/rerank), sequence classification (/predict), and sparse embeddings (/embed_sparse), with support for both HTTP and gRPC APIs.

Tokens
19.7K
Snippets
62
Records
115
Agent score
90%

What's inside Text Embeddings Inference

  1. Overview of Text Embeddings Inference (TEI)

    main

    Text Embeddings Inference (TEI) is a toolkit for efficient deployment and serving of open-source text embeddings models. It is optimized for high-performance extraction of popular models such as FlagEmbedding, Ember, GTE, and E5.

    Key Capabilities:

    • Streamlined Deployment: No model graph compilation required.
    • Efficient Resource Utilization: Small Docker images and rapid boot times suitable for serverless architectures.
    • Dynamic Batching: Uses token-based dynamic batching to optimize resource utilization.
    • Optimized Inference: Leverages Flash Attention, Candle, and cuBLASLt via optimized transformers code.
    • Fast Loading: Uses Safetensors weights for faster boot times.
    • Production Features: Supports Open Telemetry for distributed tracing and exports Prometheus metrics.
  2. Supported embedding models

    main

    Text Embeddings Inference supports a wide range of embedding models across different architectures and position types:

    • Absolute positions: Nomic, BERT, CamemBERT, XLM-RoBERTa.
    • Alibi positions: JinaBERT.
    • Rope positions: Mistral, Alibaba GTE, Qwen2, MPNet, ModernBERT, Qwen3, and Gemma3.

    Commonly used models include Qwen/Qwen3-Embedding-8B, intfloat/multilingual-e5-large-instruct, nomic-ai/nomic-embed-text-v1.5, and jinaai/jina-embeddings-v2-base-en.

  3. Supported re-rankers and sequence classification models

    main

    The project supports re-ranking and sequence classification tasks using models with absolute positions:

    • Re-Ranking: Supports XLM-RoBERTa (e.g., BAAI/bge-reranker-large), GTE (e.g., Alibaba-NLP/gte-multilingual-reranker-base), and ModernBert (e.g., Alibaba-NLP/gte-reranker-modernbert-base).
    • Sentiment Analysis: Supports RoBERTa (e.g., SamLowe/roberta-base-go_emotions).
  4. Configure Cloud NAT for TEI on Cloud Run

    main

    If you need to route all traffic through a VPC network to reach the public internet (e.g., for improved network performance), you must set up Google Cloud NAT. This requires creating a Cloud Router and a NAT gateway in your region.

    gcloud compute routers create nat-router --network=default --region=$LOCATION
    gcloud compute routers nats create vm-nat --router=nat-router --region=$LOCATION --auto-allocate-nat-external-ips --nat-all-subnet-ip-ranges
  5. Select the appropriate Docker image for your hardware

    main

    Choose a Docker image based on your hardware architecture and GPU generation. For GPU workloads, ensure you have the NVIDIA Container Toolkit installed and NVIDIA drivers with CUDA version 12.2 or higher.

    ArchitecturePlatformImage
    CPUx86_64ghcr.io/huggingface/text-embeddings-inference:cpu-1.9
    CPUaarch64ghcr.io/huggingface/text-embeddings-inference:cpu-arm64-1.9
    Turing (T4, RTX 2000 series, ...)x86_64ghcr.io/huggingface/text-embeddings-inference:turing-1.9 (experimental)
    Ampere 8.0 (A100, A30)x86_64ghcr.io/huggingface/text-embeddings-inference:1.9
    Ampere 8.6 (A10, A40, ...)x86_64ghcr.io/huggingface/text-embeddings-inference:86-1.9
    Ada Lovelace (RTX 4000 series, ...)x86_64ghcr.io/huggingface/text-embeddings-inference:89-1.9
    Hopper (H100)x86_64ghcr.io/huggingface/text-embeddings-inference:hopper-1.9
    Blackwell 10.0 (B200, GB200, ...)x86_64ghcr.io/huggingface/text-embeddings-inference:100-1.9 (experimental)
    Blackwell 12.0 (GeForce RTX 50X0, ...)x86_64ghcr.io/huggingface/text-embeddings-inference:120-1.9 (experimental)
    Blackwell 12.1 (DGX Spark GB10, ...)multighcr.io/huggingface/text-embeddings-inference:121-1.9 (experimental)
  6. Perform Re-ranker Inference

    main

    TEI supports re-rankers (cross-encoders) via the /rerank endpoint. Deploy a re-ranker model (e.g., BAAI/bge-reranker-large) using the standard Docker command, then query it using cURL.

    curl 127.0.0.1:8080/rerank \
        -X POST \
        -d '{"query":"What is Deep Learning?", "texts": ["Deep Learning is not...", "Deep learning is..."], "raw_scores": false}' \
        -H 'Content-Type: application/json'
  7. Deploy TEI on Cloud Run (CPU only)

    main

    Deploy TEI as a serverless service on Google Cloud Run using only CPU resources. This command uses the official Docker image and sets up the service with specific concurrency and memory limits.

    gcloud run deploy $SERVICE_NAME \
        --image=$CONTAINER_URI \
        --args="--model-id=$MODEL_ID,--max-concurrent-requests=64" \
        --set-env-vars=HF_HUB_ENABLE_HF_TRANSFER=1 \
        --port=8080 \
        --cpu=8 \
        --memory=32Gi \
        --region=$LOCATION \
        --no-allow-unauthenticated
  8. Install Text Embeddings Inference locally with GPU

    main

    To run text-embeddings-inference locally on your own machine using a GPU, follow these steps to set up your environment and compile the necessary CUDA kernels.

    Prerequisites

    1. CUDA and NVIDIA Drivers: Ensure NVIDIA drivers are installed and compatible with CUDA version 12.2 or higher.
    2. Path Configuration: Add NVIDIA binaries to your PATH:
      export PATH=$PATH:/usr/local/cuda/bin
    3. Rust: Install Rust using the official installer:
      curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

    Compilation

    Compilation may take a significant amount of time due to the compilation of CUDA kernels. Choose the command corresponding to your GPU architecture:

    • For Turing GPUs (e.g., T4, RTX 2000 series):
      cargo install --path router -F candle-cuda-turing
    • For Ampere, Ada Lovelace, Hopper, and Blackwell GPUs:
      cargo install --path router -F candle-cuda
    export PATH=$PATH:/usr/local/cuda/bin
    
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
    
    # For Turing GPUs
    cargo install --path router -F candle-cuda-turing
    
    # For Ampere, Ada Lovelace, Hopper, and Blackwell
    cargo install --path router -F candle-cuda