Spaces:
Running
Running
| from openai import AsyncOpenAI | |
| import logging | |
| from typing import List, Dict, Union, Optional | |
| from datasets import Dataset | |
| import asyncio | |
| import numpy as np | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from src.api.exceptions import OpenAIError | |
| # Set up structured logging | |
| logging.basicConfig( | |
| level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" | |
| ) | |
| logger = logging.getLogger(__name__) | |
| class EmbeddingService: | |
| def __init__( | |
| self, | |
| openai_api_key: str, | |
| model: str = "text-embedding-3-small", | |
| batch_size: int = 10, | |
| max_concurrent_requests: int = 10, # Limit to 10 concurrent requests | |
| ): | |
| self.client = AsyncOpenAI(api_key=openai_api_key) | |
| self.model = model | |
| self.batch_size = batch_size | |
| self.semaphore = asyncio.Semaphore(max_concurrent_requests) # Rate limiter | |
| self.total_requests = 0 # Total number of requests to process | |
| self.completed_requests = 0 # Number of completed requests | |
| async def get_embedding(self, text: str) -> List[float]: | |
| """Generate embeddings for the given text using OpenAI.""" | |
| text = text.replace("\n", " ") | |
| try: | |
| async with self.semaphore: # Acquire a semaphore slot | |
| response = await self.client.embeddings.create( | |
| input=[text], model=self.model | |
| ) | |
| self.completed_requests += 1 # Increment completed requests | |
| self._log_progress() # Log progress | |
| return response.data[0].embedding | |
| except Exception as e: | |
| logger.error(f"Failed to generate embedding: {e}") | |
| raise OpenAIError(f"OpenAI API error: {e}") | |
| async def create_embeddings( | |
| self, | |
| data: Union[Dataset, List[str]], | |
| target_column: str = None, | |
| output_column: str = "embeddings", | |
| ) -> Union[Dataset, List[List[float]]]: | |
| """ | |
| Create embeddings for either a Dataset or a list of strings. | |
| Args: | |
| data: Either a Dataset or a list of strings. | |
| target_column: The column in the Dataset to generate embeddings for (required if data is a Dataset). | |
| output_column: The column to store embeddings in the Dataset (default: "embeddings"). | |
| Returns: | |
| If data is a Dataset, returns the Dataset with the embeddings column. | |
| If data is a list of strings, returns a list of embeddings. | |
| """ | |
| if isinstance(data, Dataset): | |
| if not target_column: | |
| raise ValueError("target_column is required when data is a Dataset.") | |
| return await self._create_embeddings_for_dataset( | |
| data, target_column, output_column | |
| ) | |
| elif isinstance(data, list): | |
| return await self._create_embeddings_for_texts(data) | |
| else: | |
| raise TypeError( | |
| "data must be either a Hugging Face Dataset or a list of strings." | |
| ) | |
| async def _create_embeddings_for_dataset( | |
| self, dataset: Dataset, target_column: str, output_column: str | |
| ) -> Dataset: | |
| """Create embeddings for the target column in the Dataset.""" | |
| logger.info("Generating embeddings for Dataset...") | |
| self.total_requests = len(dataset) # Set total number of requests | |
| self.completed_requests = 0 # Reset completed requests counter | |
| embeddings = [] | |
| for i in range(0, len(dataset), self.batch_size): | |
| batch = dataset[i : i + self.batch_size] | |
| batch_embeddings = await asyncio.gather( | |
| *[self.get_embedding(text) for text in batch[target_column]] | |
| ) | |
| embeddings.extend(batch_embeddings) | |
| dataset = dataset.add_column(output_column, embeddings) | |
| return dataset | |
| async def _create_embeddings_for_texts(self, texts: List[str]) -> List[List[float]]: | |
| """Create embeddings for a list of strings.""" | |
| logger.info("Generating embeddings for list of texts...") | |
| self.total_requests = len(texts) # Set total number of requests | |
| self.completed_requests = 0 # Reset completed requests counter | |
| batches = [ | |
| texts[i : i + self.batch_size] | |
| for i in range(0, len(texts), self.batch_size) | |
| ] | |
| embeddings = [] | |
| for batch in batches: | |
| batch_embeddings = await asyncio.gather( | |
| *[self.get_embedding(text) for text in batch] | |
| ) | |
| embeddings.extend(batch_embeddings) | |
| return embeddings | |
| def _log_progress(self): | |
| """Log the progress of embedding generation.""" | |
| progress = (self.completed_requests / self.total_requests) * 100 | |
| logger.info( | |
| f"Progress: {self.completed_requests}/{self.total_requests} ({progress:.2f}%)" | |
| ) | |
| # async def search_embeddings( | |
| # self, | |
| # query_embeddings: List[List[float]], | |
| # dataset: Dataset, | |
| # embedding_column: str, | |
| # target_column: str, | |
| # num_results: int, | |
| # ) -> Dict[str, List]: | |
| # """ | |
| # Perform a cosine similarity search between query embeddings and dataset embeddings. | |
| # Args: | |
| # query_embeddings: List of embeddings for the query texts. | |
| # dataset: The dataset to search in. | |
| # embedding_column: The column in the dataset containing embeddings. | |
| # target_column: The column to return in the results. | |
| # num_results: The number of results to return. | |
| # Returns: | |
| # A dictionary of lists containing the target column values and their similarity scores. | |
| # """ | |
| # dataset_embeddings = np.array(dataset[embedding_column]) | |
| # query_embeddings = np.array(query_embeddings) | |
| # # Compute cosine similarity | |
| # similarities = cosine_similarity(query_embeddings, dataset_embeddings) | |
| # # Initialize the results dictionary | |
| # results = { | |
| # target_column: [], | |
| # "similarity": [], | |
| # } | |
| # # Get the top-k results for each query | |
| # for query_similarities in similarities: | |
| # top_k_indices = np.argsort(query_similarities)[-num_results:][::-1] | |
| # for idx in top_k_indices: | |
| # results[target_column].append(dataset[target_column][idx]) | |
| # results["similarity"].append(float(query_similarities[idx])) | |
| # return results | |
| async def search_embeddings( | |
| self, | |
| query_embeddings: List[List[float]], | |
| dataset: Dataset, | |
| embedding_column: str, | |
| target_column: str, | |
| num_results: int, | |
| additional_columns: Optional[List[str]] = None, | |
| ) -> Dict[str, List]: | |
| """ | |
| Perform a cosine similarity search between query embeddings and dataset embeddings. | |
| Args: | |
| query_embeddings: List of embeddings for the query texts. | |
| dataset: The dataset to search in. | |
| embedding_column: The column in the dataset containing embeddings. | |
| target_column: The column to return in the results. | |
| num_results: The number of results to return. | |
| additional_columns: List of additional columns to include in the results. | |
| Returns: | |
| A dictionary of lists containing the target column values, their similarity scores, | |
| and any additional columns specified. | |
| """ | |
| dataset_embeddings = np.array(dataset[embedding_column]) | |
| query_embeddings = np.array(query_embeddings) | |
| # Compute cosine similarity | |
| similarities = cosine_similarity(query_embeddings, dataset_embeddings) | |
| # Initialize the results dictionary | |
| results = { | |
| target_column: [], | |
| "similarity": [], | |
| } | |
| # Add additional columns to the results dictionary | |
| if additional_columns: | |
| for column in additional_columns: | |
| results[column] = [] | |
| # Get the top-k results for each query | |
| for query_similarities in similarities: | |
| top_k_indices = np.argsort(query_similarities)[-num_results:][::-1] | |
| for idx in top_k_indices: | |
| results[target_column].append(dataset[target_column][idx]) | |
| results["similarity"].append(float(query_similarities[idx])) | |
| if additional_columns: | |
| for column in additional_columns: | |
| results[column].append(dataset[column][idx]) | |
| return results | |