tessl/pypi-flagembedding

FlagEmbedding - BGE: One-Stop Retrieval Toolkit For Search and RAG

—

Pending

Quality

Pending

Does it follow best practices?

Impact

Pending

No eval scenarios have been run

Overview

Eval results

Files

Auto Models

Name: tessl/pypi-flagembedding
Author: tessl

Auto model factory classes that automatically select and initialize the appropriate embedder or reranker implementation based on the model name. These provide the simplest way to use FlagEmbedding with any supported model without needing to know the specific model architecture or implementation details.

Capabilities

FlagAutoModel Factory

Automatically selects the appropriate embedder class based on the model name and returns an initialized instance. Supports all embedder model types including encoder-only, decoder-only, and specialized variants.

from typing import Union

class FlagAutoModel:
    @classmethod
    def from_finetuned(
        cls,
        model_name_or_path: str,
        model_class: Optional[Union[str, EmbedderModelClass]] = None,
        normalize_embeddings: bool = True,
        use_fp16: bool = True,
        query_instruction_for_retrieval: Optional[str] = None,
        devices: Optional[Union[str, List[str]]] = None,
        pooling_method: Optional[str] = None,
        trust_remote_code: Optional[bool] = None,
        query_instruction_format: Optional[str] = None,
        **kwargs
    ) -> AbsEmbedder:
        """
        Automatically create embedder instance from model name.
        
        Args:
            model_name_or_path: Path to model or HuggingFace model name
            model_class: Force specific model class (optional)
            normalize_embeddings: Whether to normalize output embeddings
            use_fp16: Use half precision for inference
            query_instruction_for_retrieval: Instruction prepended to queries
            devices: List of devices for multi-GPU inference
            pooling_method: Pooling strategy ("cls", "mean", "last_token")
            trust_remote_code: Allow custom model code execution
            query_instruction_format: Format string for query instructions
            **kwargs: Additional model-specific parameters
            
        Returns:
            Initialized AbsEmbedder instance
        """

FlagAutoReranker Factory

Automatically selects the appropriate reranker class based on the model name and returns an initialized instance. Supports all reranker model types including encoder-only, decoder-only, and specialized variants.

class FlagAutoReranker:
    @classmethod
    def from_finetuned(
        cls,
        model_name_or_path: str,
        model_class: Optional[Union[str, RerankerModelClass]] = None,
        use_fp16: bool = False,
        trust_remote_code: Optional[bool] = None,
        **kwargs
    ) -> AbsReranker:
        """
        Automatically create reranker instance from model name.
        
        Args:
            model_name_or_path: Path to model or HuggingFace model name
            model_class: Force specific model class (optional)
            use_fp16: Use half precision for inference
            trust_remote_code: Allow custom model code execution
            **kwargs: Additional model-specific parameters
            
        Returns:
            Initialized AbsReranker instance
        """

Usage Examples

Basic Auto Embedder Usage

from FlagEmbedding import FlagAutoModel

# Automatic model selection
embedder = FlagAutoModel.from_finetuned('bge-large-en-v1.5')

# Encode single query
embedding = embedder.encode_queries("What is machine learning?")

# Encode multiple documents
docs = ["ML is a subset of AI", "Neural networks are computing systems"]
doc_embeddings = embedder.encode_corpus(docs)

Multi-GPU Auto Embedder

from FlagEmbedding import FlagAutoModel

# Use multiple GPUs for faster inference
embedder = FlagAutoModel.from_finetuned(
    'bge-large-en-v1.5',
    devices=['cuda:0', 'cuda:1', 'cuda:2'],
    use_fp16=True
)

# Process large batch across multiple devices
large_corpus = [f"Document {i}" for i in range(10000)]
embeddings = embedder.encode_corpus(large_corpus, batch_size=512)

Custom Instructions and Formatting

from FlagEmbedding import FlagAutoModel

# Use custom query instructions
embedder = FlagAutoModel.from_finetuned(
    'bge-large-en-v1.5',
    query_instruction_for_retrieval="Represent this query for retrieving relevant documents: ",
    query_instruction_format="{}{}",
    normalize_embeddings=True
)

queries = ["How to implement neural networks"]
embeddings = embedder.encode_queries(queries)

Auto Reranker Usage

from FlagEmbedding import FlagAutoReranker

# Initialize reranker
reranker = FlagAutoReranker.from_finetuned('bge-reranker-base')

# Score query-document pairs
pairs = [
    ("machine learning query", "ML is a subset of artificial intelligence"),
    ("machine learning query", "Cooking recipes for pasta dishes")
]

scores = reranker.compute_score(pairs)
print(f"Relevance scores: {scores}")  # Higher scores = more relevant

Forced Model Class Selection

from FlagEmbedding import FlagAutoModel, EmbedderModelClass

# Force specific model class even if auto-detection would choose different
embedder = FlagAutoModel.from_finetuned(
    'custom-model-name',
    model_class=EmbedderModelClass.ENCODER_ONLY_BASE,
    pooling_method="mean",
    trust_remote_code=True
)

Model Selection Logic

The auto factories use the following logic to select model implementations:

Model Name Matching: Checks model name against known model patterns
Architecture Detection: Analyzes model configuration to determine architecture type
Fallback Selection: Uses sensible defaults based on model characteristics
Override Support: Allows manual override via model_class parameter

Error Handling

from FlagEmbedding import FlagAutoModel

try:
    embedder = FlagAutoModel.from_finetuned('unknown-model')
except ValueError as e:
    print(f"Model not supported: {e}")
    
try:
    # Invalid device specification
    embedder = FlagAutoModel.from_finetuned('bge-base-en', devices=['invalid:0'])
except RuntimeError as e:
    print(f"Device error: {e}")

Types

from typing import Optional, List
from FlagEmbedding import AbsEmbedder, AbsReranker, EmbedderModelClass, RerankerModelClass

# Type aliases for auto model factory methods
AutoEmbedderFactory = type(FlagAutoModel.from_finetuned)
AutoRerankerFactory = type(FlagAutoReranker.from_finetuned)

Install with Tessl CLI