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

Model Types and Utilities

Name: tessl/pypi-flagembedding
Author: tessl

Enumerations for supported model classes and utility functions for discovering available models and their capabilities. These components help with model selection, validation, and programmatic access to FlagEmbedding's model ecosystem.

Capabilities

EmbedderModelClass Enumeration

Enumeration defining the available embedder model architecture classes, used for programmatic model selection and validation.

from enum import Enum

class EmbedderModelClass(Enum):
    """Enumeration of available embedder model classes."""
    
    ENCODER_ONLY_BASE = "encoder-only-base"
    """Standard encoder-only models (BERT-like architectures)."""
    
    ENCODER_ONLY_M3 = "encoder-only-m3"
    """BGE-M3 specialized encoder models with multi-vector support."""
    
    DECODER_ONLY_BASE = "decoder-only-base"
    """Standard decoder-only models (LLM-like architectures)."""
    
    DECODER_ONLY_ICL = "decoder-only-icl"
    """In-context learning decoder models."""

RerankerModelClass Enumeration

Enumeration defining the available reranker model architecture classes, used for programmatic reranker selection and validation.

from enum import Enum

class RerankerModelClass(Enum):
    """Enumeration of available reranker model classes."""
    
    ENCODER_ONLY_BASE = "encoder-only-base"
    """Standard encoder-only rerankers (cross-encoder architecture)."""
    
    DECODER_ONLY_BASE = "decoder-only-base"
    """Standard decoder-only rerankers (LLM-based)."""
    
    DECODER_ONLY_LAYERWISE = "decoder-only-layerwise"
    """Layer-wise processing decoder rerankers."""
    
    DECODER_ONLY_LIGHTWEIGHT = "decoder-only-lightweight"
    """Lightweight decoder rerankers for efficiency."""

PoolingMethod Enumeration

Enumeration of pooling strategies available for embedders, determining how token representations are combined into sentence embeddings.

from enum import Enum

class PoolingMethod(Enum):
    """Enumeration of pooling methods for embedders."""
    
    LAST_TOKEN = "last_token"
    """Use the last token representation (common for decoder-only models)."""
    
    CLS = "cls"
    """Use the CLS token representation (common for encoder-only models)."""
    
    MEAN = "mean"
    """Use mean pooling across all tokens."""

Model Discovery Utilities

Utility functions for discovering available models and their capabilities programmatically.

def support_model_list() -> List[str]:
    """
    Get list of all supported model names across all model types.
    
    Returns:
        List of all supported model names that can be used with
        FlagAutoModel and FlagAutoReranker
    """

def support_native_bge_model_list() -> List[str]:
    """
    Get list of native BGE model names specifically supported.
    
    Returns:
        List of BGE model names with optimized support
    """

Usage Examples

Programmatic Model Selection

from FlagEmbedding import EmbedderModelClass, RerankerModelClass, FlagAutoModel

# Use enumeration for type-safe model selection
model_class = EmbedderModelClass.ENCODER_ONLY_BASE

embedder = FlagAutoModel.from_finetuned(
    'bge-large-en-v1.5',
    model_class=model_class,  # Explicit model class
    use_fp16=True
)

# Check model class programmatically
if model_class == EmbedderModelClass.ENCODER_ONLY_M3:
    print("Using BGE-M3 specialized model")
elif model_class == EmbedderModelClass.DECODER_ONLY_BASE:
    print("Using LLM-based embedder")

Model Discovery and Validation

from FlagEmbedding import support_model_list, support_native_bge_model_list

# Get all supported models
all_models = support_model_list()
print(f"Total supported models: {len(all_models)}")

# Get BGE-specific models
bge_models = support_native_bge_model_list()
print(f"Native BGE models: {len(bge_models)}")

# Validate model availability
def is_model_supported(model_name: str) -> bool:
    return model_name in support_model_list()

# Check specific models
test_models = ['bge-large-en-v1.5', 'custom-model', 'e5-large-v2']
for model in test_models:
    status = "✓" if is_model_supported(model) else "✗"
    print(f"{status} {model}")

Dynamic Model Configuration

from FlagEmbedding import EmbedderModelClass, PoolingMethod, FlagAutoModel

# Configuration based on model type
def get_optimal_config(model_class: EmbedderModelClass) -> dict:
    """Get optimal configuration for different model classes."""
    
    if model_class == EmbedderModelClass.ENCODER_ONLY_BASE:
        return {
            'pooling_method': PoolingMethod.CLS.value,
            'batch_size': 256,
            'use_fp16': True
        }
    elif model_class == EmbedderModelClass.ENCODER_ONLY_M3:
        return {
            'pooling_method': PoolingMethod.CLS.value,
            'batch_size': 128,  # Smaller for M3
            'return_dense': True,
            'return_sparse': False
        }
    elif model_class == EmbedderModelClass.DECODER_ONLY_BASE:
        return {
            'pooling_method': PoolingMethod.LAST_TOKEN.value,
            'batch_size': 64,   # Smaller for LLM
            'use_fp16': True
        }
    else:
        return {}

# Apply configuration dynamically
model_class = EmbedderModelClass.ENCODER_ONLY_BASE
config = get_optimal_config(model_class)

embedder = FlagAutoModel.from_finetuned(
    'bge-large-en-v1.5',
    model_class=model_class,
    **config
)

Model Type Detection

from FlagEmbedding import support_model_list, EmbedderModelClass

def detect_model_type(model_name: str) -> EmbedderModelClass:
    """Detect model type based on model name patterns."""
    
    model_name_lower = model_name.lower()
    
    if 'bge-m3' in model_name_lower:
        return EmbedderModelClass.ENCODER_ONLY_M3
    elif any(llm_name in model_name_lower for llm_name in ['mistral', 'qwen', 'gemma']):
        return EmbedderModelClass.DECODER_ONLY_BASE
    elif 'icl' in model_name_lower:
        return EmbedderModelClass.DECODER_ONLY_ICL
    else:
        return EmbedderModelClass.ENCODER_ONLY_BASE

# Test model type detection
test_models = [
    'bge-large-en-v1.5',
    'bge-m3',
    'e5-mistral-7b-instruct',
    'bge-en-icl'
]

for model in test_models:
    detected_type = detect_model_type(model)
    print(f"{model} -> {detected_type.value}")

Reranker Type Selection

from FlagEmbedding import RerankerModelClass, FlagAutoReranker

def select_reranker_by_requirements(
    speed_priority: bool = False,
    accuracy_priority: bool = False,
    resource_constrained: bool = False
) -> RerankerModelClass:
    """Select optimal reranker class based on requirements."""
    
    if resource_constrained:
        return RerankerModelClass.DECODER_ONLY_LIGHTWEIGHT
    elif speed_priority:
        return RerankerModelClass.ENCODER_ONLY_BASE
    elif accuracy_priority:
        return RerankerModelClass.DECODER_ONLY_BASE
    else:
        return RerankerModelClass.DECODER_ONLY_LAYERWISE  # Balanced

# Example usage scenarios
scenarios = [
    {'speed_priority': True},
    {'accuracy_priority': True},
    {'resource_constrained': True},
    {}  # Default balanced
]

for i, scenario in enumerate(scenarios):
    reranker_class = select_reranker_by_requirements(**scenario)
    print(f"Scenario {i+1}: {reranker_class.value}")

Model Compatibility Checking

from FlagEmbedding import EmbedderModelClass, PoolingMethod

def check_model_compatibility(
    model_class: EmbedderModelClass,
    pooling_method: str
) -> bool:
    """Check if pooling method is compatible with model class."""
    
    compatible_combinations = {
        EmbedderModelClass.ENCODER_ONLY_BASE: [PoolingMethod.CLS.value, PoolingMethod.MEAN.value],
        EmbedderModelClass.ENCODER_ONLY_M3: [PoolingMethod.CLS.value],
        EmbedderModelClass.DECODER_ONLY_BASE: [PoolingMethod.LAST_TOKEN.value],
        EmbedderModelClass.DECODER_ONLY_ICL: [PoolingMethod.LAST_TOKEN.value]
    }
    
    return pooling_method in compatible_combinations.get(model_class, [])

# Test compatibility
test_cases = [
    (EmbedderModelClass.ENCODER_ONLY_BASE, PoolingMethod.CLS.value),
    (EmbedderModelClass.ENCODER_ONLY_BASE, PoolingMethod.LAST_TOKEN.value),  # Invalid
    (EmbedderModelClass.DECODER_ONLY_BASE, PoolingMethod.LAST_TOKEN.value),
    (EmbedderModelClass.ENCODER_ONLY_M3, PoolingMethod.MEAN.value)  # Invalid
]

for model_class, pooling in test_cases:
    compatible = check_model_compatibility(model_class, pooling)
    status = "✓" if compatible else "✗"
    print(f"{status} {model_class.value} + {pooling}")

Advanced Model Registry

from FlagEmbedding import support_model_list, EmbedderModelClass, RerankerModelClass
from typing import Dict, List

class ModelRegistry:
    """Advanced model registry with categorization and metadata."""
    
    def __init__(self):
        self.supported_models = support_model_list()
        self._build_registry()
    
    def _build_registry(self):
        """Build categorized model registry."""
        self.registry = {
            'embedders': {
                'bge': [m for m in self.supported_models if 'bge' in m.lower() and 'reranker' not in m.lower()],
                'e5': [m for m in self.supported_models if 'e5' in m.lower()],
                'gte': [m for m in self.supported_models if 'gte' in m.lower()]
            },
            'rerankers': [m for m in self.supported_models if 'reranker' in m.lower()]
        }
    
    def get_models_by_family(self, family: str) -> List[str]:
        """Get models by family (bge, e5, gte)."""
        return self.registry['embedders'].get(family, [])
    
    def get_rerankers(self) -> List[str]:
        """Get all reranker models."""
        return self.registry['rerankers']
    
    def recommend_model(self, task_type: str, performance_tier: str) -> str:
        """Recommend model based on task and performance requirements."""
        recommendations = {
            ('embedding', 'high'): 'bge-large-en-v1.5',
            ('embedding', 'medium'): 'bge-base-en-v1.5',
            ('embedding', 'fast'): 'bge-small-en-v1.5',
            ('reranking', 'high'): 'bge-reranker-large',
            ('reranking', 'medium'): 'bge-reranker-base',
            ('reranking', 'fast'): 'bge-reranker-v2.5-gemma2-lightweight'
        }
        return recommendations.get((task_type, performance_tier))

# Usage
registry = ModelRegistry()

print("BGE Models:", registry.get_models_by_family('bge'))
print("E5 Models:", registry.get_models_by_family('e5'))
print("Rerankers:", registry.get_rerankers())

# Get recommendations
embedding_model = registry.recommend_model('embedding', 'high')
reranking_model = registry.recommend_model('reranking', 'medium')
print(f"Recommended embedding model: {embedding_model}")
print(f"Recommended reranking model: {reranking_model}")

Model Categories

BGE Models (BAAI General Embedding)

English: bge-large-en-v1.5, bge-base-en-v1.5, bge-small-en-v1.5
Chinese: bge-large-zh-v1.5, bge-base-zh-v1.5, bge-small-zh-v1.5
Multilingual: bge-multilingual-gemma2
Specialized: bge-m3 (multi-vector), bge-en-icl (in-context learning)

E5 Models (Text Embeddings by Text-to-Text)

Standard: e5-large-v2, e5-base-v2, e5-small-v2
Multilingual: multilingual-e5-large, multilingual-e5-base, multilingual-e5-small
LLM-based: e5-mistral-7b-instruct

GTE Models (General Text Embeddings)

English: gte-large-en-v1.5, gte-base-en-v1.5
Chinese: gte-large-zh, gte-base-zh, gte-small-zh
LLM-based: gte-Qwen2-7B-instruct, gte-Qwen2-1.5B-instruct

Reranker Models

Cross-encoder: bge-reranker-base, bge-reranker-large
LLM-based: bge-reranker-v2-gemma, bge-reranker-v2-m3
Specialized: bge-reranker-v2-minicpm-layerwise, bge-reranker-v2.5-gemma2-lightweight

Types

from enum import Enum
from typing import List, Literal

# Model class enumerations
EmbedderClass = EmbedderModelClass
RerankerClass = RerankerModelClass
PoolingStrategy = PoolingMethod

# Model selection types
ModelName = str
ModelFamily = Literal["bge", "e5", "gte"]
TaskType = Literal["embedding", "reranking"]
PerformanceTier = Literal["high", "medium", "fast"]

# Utility function types
ModelList = List[str]
ModelValidator = Callable[[str], bool]

Install with Tessl CLI