Feat - Improved - Update Endpoint

docker-compose.yml

@@ -1,5 +1,3 @@
-version: "3.9"
-
 services:
   app:
     build:

src/api/dependency.py

@@ -0,0 +1,13 @@
+import os
+from src.api.services.embedding_service import EmbeddingService
+from src.api.services.huggingface_service import HuggingFaceService
+
+
+# Dependency to get EmbeddingService
+def get_embedding_service() -> EmbeddingService:
+    return EmbeddingService(openai_api_key=os.getenv("OPENAI_API_KEY"))
+
+
+# Dependency to get HuggingFaceService
+def get_huggingface_service() -> HuggingFaceService:
+    return HuggingFaceService()

src/api/models/embedding_models.py

@@ -17,10 +17,37 @@ class ReadEmbeddingRequest(BaseModel):
     dataset_name: str
 
 
+# class UpdateEmbeddingRequest(BaseModel):
+#     updates: Dict[str, List]  # Column name -> List of values
+#     target_column: str = "product_type"
+#     output_column: str = "embedding"
+#     model: str = "text-embedding-3-small"
+#     batch_size: int = 10
+#     max_concurrent_requests: int = 10
+#     dataset_name: str = "re-mind/product_type_embedding"
+
+from pydantic import BaseModel
+from typing import Dict, List
+
+
 class UpdateEmbeddingRequest(BaseModel):
-    dataset_name: str
-    updates: Dict[str, List]  # Column name -> List of values
+    dataset_name: str = "re-mind/product_type_embedding"
+    updates: Dict[
+        str, List
+    ]  # Dictionary of column names and their corresponding values
+    target_column: str = (
+        "product_type"  # Column in the new data to generate embeddings for
+    )
+    output_column: str = "embedding"  # Column to store the generated embeddings
 
 
 class DeleteEmbeddingRequest(BaseModel):
     dataset_name: str
+
+
+# Request model for the /embed endpoint
+class EmbedRequest(BaseModel):
+    texts: List[str]  # List of strings to generate embeddings for
+    output_column: str = (
+        "embeddings"  # Column to store embeddings (default: "embeddings")
+    )

src/api/services/embedding_service.py

@@ -1,146 +1,3 @@
-# from openai import AsyncOpenAI
-# import logging
-# from typing import List, Dict
-# import pandas as pd
-# import asyncio
-# 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 = 100,
-#     ):
-#         self.client = AsyncOpenAI(api_key=openai_api_key)
-#         self.model = model
-#         self.batch_size = batch_size
-
-#     async def get_embedding(self, text: str) -> List[float]:
-#         """Generate embeddings for the given text using OpenAI."""
-#         text = text.replace("\n", " ")
-#         try:
-#             response = await self.client.embeddings.create(
-#                 input=[text], model=self.model
-#             )
-#             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, df: pd.DataFrame, target_column: str, output_column: str
-#     ) -> pd.DataFrame:
-#         """Create embeddings for the target column in the dataset."""
-#         logger.info("Generating embeddings...")
-#         batches = [
-#             df[i : i + self.batch_size] for i in range(0, len(df), self.batch_size)
-#         ]
-#         processed_batches = await asyncio.gather(
-#             *[
-#                 self._process_batch(batch, target_column, output_column)
-#                 for batch in batches
-#             ]
-#         )
-#         return pd.concat(processed_batches)
-
-#     async def _process_batch(
-#         self, df_batch: pd.DataFrame, target_column: str, output_column: str
-#     ) -> pd.DataFrame:
-#         """Process a batch of rows to generate embeddings."""
-#         embeddings = await asyncio.gather(
-#             *[self.get_embedding(row[target_column]) for _, row in df_batch.iterrows()]
-#         )
-#         df_batch[output_column] = embeddings
-#         return df_batch
-
-# from openai import AsyncOpenAI
-# import logging
-# from typing import List, Dict
-# import pandas as pd
-# import asyncio
-# 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, df: pd.DataFrame, target_column: str, output_column: str
-#     ) -> pd.DataFrame:
-#         """Create embeddings for the target column in the dataset."""
-#         logger.info("Generating embeddings...")
-#         self.total_requests = len(df)  # Set total number of requests
-#         self.completed_requests = 0  # Reset completed requests counter
-
-#         batches = [
-#             df[i : i + self.batch_size] for i in range(0, len(df), self.batch_size)
-#         ]
-#         processed_batches = await asyncio.gather(
-#             *[
-#                 self._process_batch(batch, target_column, output_column)
-#                 for batch in batches
-#             ]
-#         )
-#         return pd.concat(processed_batches)
-
-#     async def _process_batch(
-#         self, df_batch: pd.DataFrame, target_column: str, output_column: str
-#     ) -> pd.DataFrame:
-#         """Process a batch of rows to generate embeddings."""
-#         embeddings = await asyncio.gather(
-#             *[self.get_embedding(row[target_column]) for _, row in df_batch.iterrows()]
-#         )
-#         df_batch[output_column] = embeddings
-#         return df_batch
-
-#     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}%)"
-#         )
-
 from openai import AsyncOpenAI
 import logging
 from typing import List, Dict, Union

src/api/services/huggingface_service.py

@@ -1,8 +1,9 @@
-from datasets import Dataset, load_dataset
+from datasets import Dataset, load_dataset, concatenate_datasets
 from huggingface_hub import HfApi, HfFolder
 import logging
 from typing import Optional, Dict, List
 import pandas as pd
+from src.api.dependency import get_embedding_service, get_huggingface_service
 from src.api.exceptions import (
     DatasetNotFoundError,
     DatasetPushError,
@@ -45,19 +46,69 @@ class HuggingFaceService:
             logger.error(f"Failed to read dataset: {e}")
             raise DatasetNotFoundError(f"Dataset not found: {e}")
 
+    # async def update_dataset(
+    #     self, dataset_name: str, updates: Dict[str, List]
+    # ) -> Optional[pd.DataFrame]:
+    #     """Update a dataset on Hugging Face Hub."""
+
+    #     embedding_service = get_embedding_service()
+
+    #     try:
+    #         df_src = await self.read_dataset(dataset_name)
+    #         df_src = Dataset.from_dict(df_src)
+    #         df_update = Dataset.from_dict(updates)
+
+    #         df = concatenate_datasets(df_src, df_update)
+
+    #         # for column, values in updates.items():
+    #         #     if column in df.columns:
+    #         #         df[column] = values
+    #         #     else:
+    #         #         logger.warning(f"Column '{column}' not found in dataset.")
+    #         # await self.push_to_hub(df, dataset_name)
+    #         # return df
+    #     except Exception as e:
+    #         logger.error(f"Failed to update dataset: {e}")
+    #         raise DatasetPushError(f"Failed to update dataset: {e}")
+
     async def update_dataset(
-        self, dataset_name: str, updates: Dict[str, List]
+        self,
+        dataset_name: str,
+        updates: Dict[str, List],
+        target_column: str,
+        output_column: str = "embeddings",
     ) -> Optional[pd.DataFrame]:
-        """Update a dataset on Hugging Face Hub."""
+        """Update a dataset on Hugging Face Hub by generating embeddings for new data and concatenating it with the existing dataset."""
         try:
-            df = await self.read_dataset(dataset_name)
-            for column, values in updates.items():
-                if column in df.columns:
-                    df[column] = values
-                else:
-                    logger.warning(f"Column '{column}' not found in dataset.")
-            await self.push_to_hub(df, dataset_name)
-            return df
+            # Step 1: Load the existing dataset from Hugging Face Hub
+            logger.info(
+                f"Loading existing dataset from Hugging Face Hub: {dataset_name}..."
+            )
+            existing_ds = await self.read_dataset(dataset_name)
+            existing_df = pd.DataFrame(existing_ds)
+
+            # Step 2: Convert the new updates into a DataFrame
+            logger.info("Converting updates to DataFrame...")
+            new_df = pd.DataFrame(updates)
+
+            # Step 3: Generate embeddings for the new data
+            logger.info("Generating embeddings for the new data...")
+            embedding_service = get_embedding_service()  # Get the embedding service
+            new_df = await embedding_service.create_embeddings(
+                new_df, target_column, output_column
+            )
+
+            # Step 4: Concatenate the existing DataFrame with the new DataFrame
+            logger.info("Concatenating existing dataset with new data...")
+            updated_df = pd.concat([existing_df, new_df], ignore_index=True)
+
+            # Step 5: Push the updated dataset back to Hugging Face Hub
+            logger.info(
+                f"Pushing updated dataset to Hugging Face Hub: {dataset_name}..."
+            )
+            await self.push_to_hub(updated_df, dataset_name)
+
+            # return updated_df
         except Exception as e:
             logger.error(f"Failed to update dataset: {e}")
             raise DatasetPushError(f"Failed to update dataset: {e}")

src/main.py

@@ -197,11 +197,13 @@ from src.api.models.embedding_models import (
     ReadEmbeddingRequest,
     UpdateEmbeddingRequest,
     DeleteEmbeddingRequest,
+    EmbedRequest,
 )
 from src.api.database import get_db, Database, QueryExecutionError, HealthCheckError
 from src.api.services.embedding_service import EmbeddingService
 from src.api.services.huggingface_service import HuggingFaceService
 from src.api.exceptions import DatasetNotFoundError, DatasetPushError, OpenAIError
+from src.api.dependency import get_embedding_service, get_huggingface_service
 import pandas as pd
 import logging
 from dotenv import load_dotenv
@@ -249,24 +251,6 @@ async def health_check(db: Database = Depends(get_db)):
         raise HTTPException(status_code=500, detail=str(e))
 
 
-# Dependency to get EmbeddingService
-def get_embedding_service() -> EmbeddingService:
-    return EmbeddingService(openai_api_key=os.getenv("OPENAI_API_KEY"))
-
-
-# Dependency to get HuggingFaceService
-def get_huggingface_service() -> HuggingFaceService:
-    return HuggingFaceService()
-
-
-# Request model for the /embed endpoint
-class EmbedRequest(BaseModel):
-    texts: List[str]  # List of strings to generate embeddings for
-    output_column: str = (
-        "embeddings"  # Column to store embeddings (default: "embeddings")
-    )
-
-
 # Endpoint to generate embeddings for a list of strings
 @app.post("/embed")
 async def embed(
@@ -363,21 +347,51 @@ async def read_embeddings(
 
 
 # Endpoint to update embeddings
+# @app.post("/update_embeddings")
+# async def update_embeddings(
+#     request: UpdateEmbeddingRequest,
+#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
+# ):
+#     """
+#     Update embeddings in a Hugging Face dataset.
+#     """
+#     try:
+#         df = await huggingface_service.update_dataset(
+#             request.dataset_name, request.updates
+#         )
+#         return {
+#             "message": "Embeddings updated successfully.",
+#             "dataset_name": request.dataset_name,
+#         }
+#     except DatasetPushError as e:
+#         logger.error(f"Failed to update dataset: {e}")
+#         raise HTTPException(status_code=500, detail=f"Failed to update dataset: {e}")
+#     except Exception as e:
+#         logger.error(f"An error occurred: {e}")
+#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
 @app.post("/update_embeddings")
 async def update_embeddings(
     request: UpdateEmbeddingRequest,
     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
 ):
     """
-    Update embeddings in a Hugging Face dataset.
+    Update embeddings in a Hugging Face dataset by generating embeddings for new data and concatenating it with the existing dataset.
     """
     try:
-        df = await huggingface_service.update_dataset(
-            request.dataset_name, request.updates
+        # Call the update_dataset method to generate embeddings, concatenate, and push the updated dataset
+        updated_df = await huggingface_service.update_dataset(
+            request.dataset_name,
+            request.updates,
+            request.target_column,
+            request.output_column,
         )
+
         return {
             "message": "Embeddings updated successfully.",
             "dataset_name": request.dataset_name,
+            "num_rows": len(updated_df),
         }
     except DatasetPushError as e:
         logger.error(f"Failed to update dataset: {e}")