embeddings.py

import torch
from langchain.embeddings import HuggingFaceEmbeddings


def get_hf_embeddings(model_name=None):
    """Get huggingface embedding."""
    
    if model_name is None:
        # Some candiates: 
        # "BAAI/bge-m3" (good, though large and slow)
        # "BAAI/bge-base-en-v1.5" -> seems not that good with current settings
        # "sentence-transformers/all-mpnet-base-v2", "maidalun1020/bce-embedding-base_v1", "intfloat/multilingual-e5-large"
        # Ref: https://huggingface.co/spaces/mteb/leaderboard, https://huggingface.co/maidalun1020/bce-embedding-base_v1
        model_name = "BAAI/bge-large-en-v1.5"  # or ""
    
    embeddings = HuggingFaceEmbeddings(model_name=model_name)
    
    return embeddings


def get_jinaai_embeddings(model_name="jinaai/jina-embeddings-v2-base-en", device="auto"):
    """Get jinaai embedding."""
    
    # device: cpu or cuda
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    # For jinaai. Ref: https://github.com/langchain-ai/langchain/issues/6080
    from transformers import AutoModel
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)  # -> will yield error, need bug fixing

    model_name = model_name
    model_kwargs = {'device': device, "trust_remote_code": True}
    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
    )
    
    return embeddings