Spaces:
Sleeping
Sleeping
File size: 2,492 Bytes
b27eb78 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
from typing import Literal, Union
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import OpenAIEmbeddings
EmbeddingModelType = Literal["openai", "huggingface"]
EmbeddingsModel = Union[OpenAIEmbeddings, HuggingFaceEmbeddings]
def get_embedding_model(
model_id: str,
model_type: EmbeddingModelType = "huggingface",
device: str = "cpu",
) -> EmbeddingsModel:
"""Gets an instance of the configured embedding model.
The function returns either an OpenAI or HuggingFace embedding model based on the
provided model type.
Args:
model_id (str): The ID/name of the embedding model to use
model_type (EmbeddingModelType): The type of embedding model to use.
Must be either "openai" or "huggingface". Defaults to "huggingface"
device (str): The device to use for the embedding model. Defaults to "cpu"
Returns:
EmbeddingsModel: An embedding model instance based on the configuration settings
Raises:
ValueError: If model_type is not "openai" or "huggingface"
"""
if model_type == "openai":
return get_openai_embedding_model(model_id)
elif model_type == "huggingface":
return get_huggingface_embedding_model(model_id, device)
else:
raise ValueError(f"Invalid embedding model type: {model_type}")
def get_openai_embedding_model(model_id: str) -> OpenAIEmbeddings:
"""Gets an OpenAI embedding model instance.
Args:
model_id (str): The ID/name of the OpenAI embedding model to use
Returns:
OpenAIEmbeddings: A configured OpenAI embeddings model instance with
special token handling enabled
"""
return OpenAIEmbeddings(
model=model_id,
allowed_special={"<|endoftext|>"},
)
def get_huggingface_embedding_model(
model_id: str, device: str
) -> HuggingFaceEmbeddings:
"""Gets a HuggingFace embedding model instance.
Args:
model_id (str): The ID/name of the HuggingFace embedding model to use
device (str): The compute device to run the model on (e.g. "cpu", "cuda")
Returns:
HuggingFaceEmbeddings: A configured HuggingFace embeddings model instance
with remote code trust enabled and embedding normalization disabled
"""
return HuggingFaceEmbeddings(
model_name=model_id,
model_kwargs={"device": device, "trust_remote_code": True},
encode_kwargs={"normalize_embeddings": False},
)
|