Views
No views yet
1from typing import List
2from llama_index.embeddings.huggingface_optimum import OptimumEmbedding
3import asyncio
4
5class CustomEmbedding:
6 def __init__(self, folder_name: str):
7 """Initialize the embedding model."""
8 self.embed_model = OptimumEmbedding(folder_name=folder_name)
9
10 async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
11 """Asynchronously embed a list of documents."""
12 loop = asyncio.get_event_loop()
13 return await loop.run_in_executor(None, self.embed_documents, texts)
14
15 async def aembed_query(self, text: str) -> List[float]:
16 """Asynchronously embed a single query."""
17 loop = asyncio.get_event_loop()
18 return await loop.run_in_executor(None, self.embed_query, text)
19
20 def embed_documents(self, texts: List[str]) -> List[List[float]]:
21 """Embed a list of documents."""
22 return [self.embed_model.get_text_embedding(text) for text in texts]
23
24 def embed_query(self, text: str) -> List[float]:
25 """Embed a single query."""
26 return self.embed_model.get_text_embedding(text)
27
28# Example Usage
29custom_embeddings = CustomEmbedding(folder_name="./optimum_model")CustomEmbedding class initializes the OptimumEmbedding instance with the specified folder_name for the preloaded model.aembed_documents(texts: List[str]): Asynchronously embeds a list of documents and returns their embeddings.aembed_query(text: str): Asynchronously embeds a single query and returns its embedding.embed_documents(texts: List[str]): Embeds a list of documents and returns their embeddings.embed_query(text: str): Embeds a single query and returns its embedding."./optimum_model" with the path to your locally stored Optimum ONNX Runtime model.1# Embed a single query
2query_embedding = custom_embeddings.embed_query("Hello World!")
3
4# Embed multiple documents
5document_embeddings = custom_embeddings.embed_documents(["Document 1", "Document 2"])