llama-index

1
from typing import Dict, Type
2

3
from llama_index.legacy.embeddings.base import BaseEmbedding
4
from llama_index.legacy.embeddings.google import GoogleUnivSentEncoderEmbedding
5
from llama_index.legacy.embeddings.huggingface import HuggingFaceEmbedding
6
from llama_index.legacy.embeddings.langchain import LangchainEmbedding
7
from llama_index.legacy.embeddings.openai import OpenAIEmbedding
8
from llama_index.legacy.embeddings.text_embeddings_inference import (
9
    TextEmbeddingsInference,
10
)
11
from llama_index.legacy.embeddings.utils import resolve_embed_model
12
from llama_index.legacy.token_counter.mock_embed_model import MockEmbedding
13

14
RECOGNIZED_EMBEDDINGS: Dict[str, Type[BaseEmbedding]] = {
15
    GoogleUnivSentEncoderEmbedding.class_name(): GoogleUnivSentEncoderEmbedding,
16
    OpenAIEmbedding.class_name(): OpenAIEmbedding,
17
    LangchainEmbedding.class_name(): LangchainEmbedding,
18
    MockEmbedding.class_name(): MockEmbedding,
19
    HuggingFaceEmbedding.class_name(): HuggingFaceEmbedding,
20
    TextEmbeddingsInference.class_name(): TextEmbeddingsInference,
21
    OpenAIEmbedding.class_name(): OpenAIEmbedding,
22
}
23

24

25
def load_embed_model(data: dict) -> BaseEmbedding:
26
    """Load Embedding by name."""
27
    if isinstance(data, BaseEmbedding):
28
        return data
29
    name = data.get("class_name", None)
30
    if name is None:
31
        raise ValueError("Embedding loading requires a class_name")
32
    if name not in RECOGNIZED_EMBEDDINGS:
33
        raise ValueError(f"Invalid Embedding name: {name}")
34

35
    # special handling for LangchainEmbedding
36
    # it can be any local model technially
37
    if name == LangchainEmbedding.class_name():
38
        local_name = data.get("model_name", None)
39
        if local_name is not None:
40
            return resolve_embed_model("local:" + local_name)
41
        else:
42
            raise ValueError("LangchainEmbedding requires a model_name")
43

44
    return RECOGNIZED_EMBEDDINGS[name].from_dict(data)
45