"""
Load embedding models from huggingface.
"""

import torch
from langchain.embeddings import HuggingFaceEmbeddings


def get_hf_embeddings(model_name=None):
    """Get huggingface embedding by name."""

    if model_name is None:
        # Some candiates
        # "BAAI/bge-m3" (good, though large and slow)
        # "BAAI/bge-base-en-v1.5" -> also good
        # "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"

    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