rokeya71/granite-embedding-30m-english-optimum-q8

Converted and quantized.

Custom Embedding Class for Optimum ONNX Runtime

This document provides the implementation of a custom embedding class designed to work with the Optimum ONNX Runtime model.

from typing import List
from llama_index.embeddings.huggingface_optimum import OptimumEmbedding
import asyncio

class CustomEmbedding:
    def __init__(self, folder_name: str):
        """Initialize the embedding model."""
        self.embed_model = OptimumEmbedding(folder_name=folder_name)

    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """Asynchronously embed a list of documents."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.embed_documents, texts)

    async def aembed_query(self, text: str) -> List[float]:
        """Asynchronously embed a single query."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self.embed_query, text)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed a list of documents."""
        return [self.embed_model.get_text_embedding(text) for text in texts]

    def embed_query(self, text: str) -> List[float]:
        """Embed a single query."""
        return self.embed_model.get_text_embedding(text)

# Example Usage
custom_embeddings = CustomEmbedding(folder_name="./optimum_model")

Key Features

1. Initialization:

   • The CustomEmbedding class initializes the OptimumEmbedding instance with the specified folder_name for the preloaded model.

2. Asynchronous Methods:

   • 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.

3. Synchronous Methods:

   • 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.

Usage

• Folder Name: Replace "./optimum_model" with the path to your locally stored Optimum ONNX Runtime model.

• Example:

  # Embed a single query
  query_embedding = custom_embeddings.embed_query("Hello World!")
  
  # Embed multiple documents
  document_embeddings = custom_embeddings.embed_documents(["Document 1", "Document 2"])