Feat - Use huggingface dataset instead of pandas

src/api/services/embedding_service.py

@@ -1,7 +1,136 @@
+# from openai import AsyncOpenAI
+# import logging
+# from typing import List, Dict, Union
+# 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,
+#         data: Union[pd.DataFrame, List[str]],
+#         target_column: str = None,
+#         output_column: str = "embeddings",
+#     ) -> Union[pd.DataFrame, List[List[float]]]:
+#         """
+#         Create embeddings for either a DataFrame or a list of strings.
+
+#         Args:
+#             data: Either a DataFrame or a list of strings.
+#             target_column: The column in the DataFrame to generate embeddings for (required if data is a DataFrame).
+#             output_column: The column to store embeddings in the DataFrame (default: "embeddings").
+
+#         Returns:
+#             If data is a DataFrame, returns the DataFrame with the embeddings column.
+#             If data is a list of strings, returns a list of embeddings.
+#         """
+#         if isinstance(data, pd.DataFrame):
+#             if not target_column:
+#                 raise ValueError("target_column is required when data is a DataFrame.")
+#             return await self._create_embeddings_for_dataframe(
+#                 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 pandas DataFrame or a list of strings."
+#             )
+
+#     async def _create_embeddings_for_dataframe(
+#         self, df: pd.DataFrame, target_column: str, output_column: str
+#     ) -> pd.DataFrame:
+#         """Create embeddings for the target column in the DataFrame."""
+#         logger.info("Generating embeddings for DataFrame...")
+#         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 _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
+
+#     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
-import pandas as pd
+from datasets import Dataset
 import asyncio
 from src.api.exceptions import OpenAIError
 
@@ -44,53 +173,53 @@ class EmbeddingService:
 
     async def create_embeddings(
         self,
-        data: Union[pd.DataFrame, List[str]],
+        data: Union[Dataset, List[str]],
         target_column: str = None,
         output_column: str = "embeddings",
-    ) -> Union[pd.DataFrame, List[List[float]]]:
+    ) -> Union[Dataset, List[List[float]]]:
         """
-        Create embeddings for either a DataFrame or a list of strings.
+        Create embeddings for either a Dataset or a list of strings.
 
         Args:
-            data: Either a DataFrame or a list of strings.
-            target_column: The column in the DataFrame to generate embeddings for (required if data is a DataFrame).
-            output_column: The column to store embeddings in the DataFrame (default: "embeddings").
+            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 DataFrame, returns the DataFrame with the embeddings column.
+            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, pd.DataFrame):
+        if isinstance(data, Dataset):
             if not target_column:
-                raise ValueError("target_column is required when data is a DataFrame.")
-            return await self._create_embeddings_for_dataframe(
+                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 pandas DataFrame or a list of strings."
+                "data must be either a Hugging Face Dataset or a list of strings."
             )
 
-    async def _create_embeddings_for_dataframe(
-        self, df: pd.DataFrame, target_column: str, output_column: str
-    ) -> pd.DataFrame:
-        """Create embeddings for the target column in the DataFrame."""
-        logger.info("Generating embeddings for DataFrame...")
-        self.total_requests = len(df)  # Set total number of requests
+    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
 
-        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)
+        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."""
@@ -110,16 +239,6 @@ class EmbeddingService:
             embeddings.extend(batch_embeddings)
         return embeddings
 
-    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

src/api/services/huggingface_service.py

@@ -1,9 +1,111 @@
+# from datasets import Dataset, load_dataset, concatenate_datasets
+# from huggingface_hub import HfApi, HfFolder
+# import logging
+# import os
+# from typing import Optional, Dict, List
+# import pandas as pd
+# from src.api.services.embedding_service import EmbeddingService
+# from src.api.exceptions import (
+#     DatasetNotFoundError,
+#     DatasetPushError,
+#     DatasetDeleteError,
+# )
+
+# # Set up structured logging
+# logging.basicConfig(
+#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+# )
+# logger = logging.getLogger(__name__)
+
+
+# class HuggingFaceService:
+#     def __init__(self, hf_token: Optional[str] = None):
+#         """Initialize the HuggingFaceService with an optional token."""
+#         self.hf_api = HfApi()
+#         if hf_token:
+#             HfFolder.save_token(hf_token)  # Save the token for authentication
+
+#     async def push_to_hub(self, df: pd.DataFrame, dataset_name: str) -> None:
+#         """Push the dataset to Hugging Face Hub."""
+#         try:
+#             logger.info(f"Creating Hugging Face Dataset: {dataset_name}...")
+#             ds = Dataset.from_pandas(df)
+#             ds.push_to_hub(dataset_name)
+#             logger.info(f"Dataset pushed to Hugging Face Hub: {dataset_name}")
+#         except Exception as e:
+#             logger.error(f"Failed to push dataset to Hugging Face Hub: {e}")
+#             raise DatasetPushError(f"Failed to push dataset: {e}")
+
+#     async def read_dataset(self, dataset_name: str) -> Optional[pd.DataFrame]:
+#         """Read a dataset from Hugging Face Hub."""
+#         try:
+#             logger.info(f"Loading dataset from Hugging Face Hub: {dataset_name}...")
+#             ds = load_dataset(dataset_name)
+#             df = ds["train"].to_dict()
+#             return df
+#         except Exception as e:
+#             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],
+#         target_column: str,
+#         output_column: str = "embeddings",
+#     ) -> Optional[pd.DataFrame]:
+#         """Update a dataset on Hugging Face Hub by generating embeddings for new data and concatenating it with the existing dataset."""
+#         try:
+#             # 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 = EmbeddingService(
+#                 openai_api_key=os.getenv("OPENAI_API_KEY")
+#             )  # 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}")
+
+#     async def delete_dataset(self, dataset_name: str) -> None:
+#         """Delete a dataset from Hugging Face Hub."""
+#         try:
+#             logger.info(f"Deleting dataset from Hugging Face Hub: {dataset_name}...")
+#             self.hf_api.delete_repo(repo_id=dataset_name, repo_type="dataset")
+#             logger.info(f"Dataset deleted from Hugging Face Hub: {dataset_name}")
+#         except Exception as e:
+#             logger.error(f"Failed to delete dataset: {e}")
+#             raise DatasetDeleteError(f"Failed to delete dataset: {e}")
+
 from datasets import Dataset, load_dataset, concatenate_datasets
 from huggingface_hub import HfApi, HfFolder
 import logging
 import os
 from typing import Optional, Dict, List
-import pandas as pd
 from src.api.services.embedding_service import EmbeddingService
 from src.api.exceptions import (
     DatasetNotFoundError,
@@ -25,24 +127,22 @@ class HuggingFaceService:
         if hf_token:
             HfFolder.save_token(hf_token)  # Save the token for authentication
 
-    async def push_to_hub(self, df: pd.DataFrame, dataset_name: str) -> None:
+    async def push_to_hub(self, dataset: Dataset, dataset_name: str) -> None:
         """Push the dataset to Hugging Face Hub."""
         try:
             logger.info(f"Creating Hugging Face Dataset: {dataset_name}...")
-            ds = Dataset.from_pandas(df)
-            ds.push_to_hub(dataset_name)
+            dataset.push_to_hub(dataset_name)
             logger.info(f"Dataset pushed to Hugging Face Hub: {dataset_name}")
         except Exception as e:
             logger.error(f"Failed to push dataset to Hugging Face Hub: {e}")
             raise DatasetPushError(f"Failed to push dataset: {e}")
 
-    async def read_dataset(self, dataset_name: str) -> Optional[pd.DataFrame]:
+    async def read_dataset(self, dataset_name: str) -> Optional[Dataset]:
         """Read a dataset from Hugging Face Hub."""
         try:
             logger.info(f"Loading dataset from Hugging Face Hub: {dataset_name}...")
-            ds = load_dataset(dataset_name)
-            df = ds["train"].to_dict()
-            return df
+            dataset = load_dataset(dataset_name)
+            return dataset["train"]
         except Exception as e:
             logger.error(f"Failed to read dataset: {e}")
             raise DatasetNotFoundError(f"Dataset not found: {e}")
@@ -53,40 +153,39 @@ class HuggingFaceService:
         updates: Dict[str, List],
         target_column: str,
         output_column: str = "embeddings",
-    ) -> Optional[pd.DataFrame]:
+    ) -> Optional[Dataset]:
         """Update a dataset on Hugging Face Hub by generating embeddings for new data and concatenating it with the existing dataset."""
         try:
             # 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)
+            existing_dataset = await self.read_dataset(dataset_name)
 
-            # Step 2: Convert the new updates into a DataFrame
-            logger.info("Converting updates to DataFrame...")
-            new_df = pd.DataFrame(updates)
+            # Step 2: Convert the new updates into a Dataset
+            logger.info("Converting updates to Dataset...")
+            new_dataset = Dataset.from_dict(updates)
 
             # Step 3: Generate embeddings for the new data
             logger.info("Generating embeddings for the new data...")
             embedding_service = EmbeddingService(
                 openai_api_key=os.getenv("OPENAI_API_KEY")
             )  # Get the embedding service
-            new_df = await embedding_service.create_embeddings(
-                new_df, target_column, output_column
+            new_dataset = await embedding_service.create_embeddings(
+                new_dataset, target_column, output_column
             )
 
-            # Step 4: Concatenate the existing DataFrame with the new DataFrame
+            # Step 4: Concatenate the existing Dataset with the new Dataset
             logger.info("Concatenating existing dataset with new data...")
-            updated_df = pd.concat([existing_df, new_df], ignore_index=True)
+            updated_dataset = concatenate_datasets([existing_dataset, new_dataset])
 
             # 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)
+            await self.push_to_hub(updated_dataset, dataset_name)
 
-            return updated_df
+            return updated_dataset
         except Exception as e:
             logger.error(f"Failed to update dataset: {e}")
             raise DatasetPushError(f"Failed to update dataset: {e}")

src/main.py

@@ -1,9 +1,259 @@
+# import os
+# from fastapi import FastAPI, Depends, HTTPException
+# from fastapi.responses import JSONResponse, RedirectResponse
+# from fastapi.middleware.gzip import GZipMiddleware
+# from pydantic import BaseModel
+# from typing import List, Dict
+# from src.api.models.embedding_models import (
+#     CreateEmbeddingRequest,
+#     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
+
+# # Load environment variables
+# load_dotenv()
+
+# # Set up structured logging
+# logging.basicConfig(
+#     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+# )
+# logger = logging.getLogger(__name__)
+
+# description = """A FastAPI application for similarity search with PostgreSQL and OpenAI embeddings.
+
+# Direct/API URL:
+# https://re-mind-similarity-search.hf.space
+# """
+
+# # Initialize FastAPI app
+# app = FastAPI(
+#     title="Similarity Search API",
+#     description=description,
+#     version="1.0.0",
+# )
+
+# app.add_middleware(GZipMiddleware, minimum_size=1000)
+
+
+# # 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()
+
+
+# # Root endpoint redirects to /docs
+# @app.get("/")
+# async def root():
+#     return RedirectResponse(url="/docs")
+
+
+# # Health check endpoint
+# @app.get("/health")
+# async def health_check(db: Database = Depends(get_db)):
+#     try:
+#         is_healthy = await db.health_check()
+#         if not is_healthy:
+#             raise HTTPException(status_code=500, detail="Database is unhealthy")
+#         return {"status": "healthy"}
+#     except HealthCheckError as e:
+#         raise HTTPException(status_code=500, detail=str(e))
+
+
+# # Endpoint to generate embeddings for a list of strings
+# @app.post("/embed")
+# async def embed(
+#     request: EmbedRequest,
+#     embedding_service: EmbeddingService = Depends(get_embedding_service),
+# ):
+#     """
+#     Generate embeddings for a list of strings and return them in the response.
+#     """
+#     try:
+#         # Step 1: Generate embeddings
+#         logger.info("Generating embeddings for list of texts...")
+#         embeddings = await embedding_service.create_embeddings(request.texts)
+
+#         return JSONResponse(
+#             content={
+#                 "message": "Embeddings generated successfully.",
+#                 "embeddings": embeddings,
+#                 "num_texts": len(request.texts),
+#             }
+#         )
+#     except OpenAIError as e:
+#         logger.error(f"OpenAI API error: {e}")
+#         raise HTTPException(status_code=500, detail=f"OpenAI API error: {e}")
+#     except Exception as e:
+#         logger.error(f"An error occurred: {e}")
+#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
+# # Endpoint to create embeddings from a database query
+# @app.post("/create_embedding")
+# async def create_embedding(
+#     request: CreateEmbeddingRequest,
+#     db: Database = Depends(get_db),
+#     embedding_service: EmbeddingService = Depends(get_embedding_service),
+#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
+# ):
+#     """
+#     Create embeddings for the target column in the dataset.
+#     """
+#     try:
+#         # Step 1: Query the database
+#         logger.info("Fetching data from the database...")
+#         result = await db.fetch(request.query)
+#         df = pd.DataFrame(result)
+
+#         # Step 2: Generate embeddings
+#         df = await embedding_service.create_embeddings(
+#             df, request.target_column, request.output_column
+#         )
+
+#         # Step 3: Push to Hugging Face Hub
+#         await huggingface_service.push_to_hub(df, request.dataset_name)
+
+#         return JSONResponse(
+#             content={
+#                 "message": "Embeddings created and pushed to Hugging Face Hub.",
+#                 "dataset_name": request.dataset_name,
+#                 "num_rows": len(df),
+#             }
+#         )
+#     except QueryExecutionError as e:
+#         logger.error(f"Database query failed: {e}")
+#         raise HTTPException(status_code=500, detail=f"Database query failed: {e}")
+#     except OpenAIError as e:
+#         logger.error(f"OpenAI API error: {e}")
+#         raise HTTPException(status_code=500, detail=f"OpenAI API error: {e}")
+#     except DatasetPushError as e:
+#         logger.error(f"Failed to push dataset: {e}")
+#         raise HTTPException(status_code=500, detail=f"Failed to push dataset: {e}")
+#     except Exception as e:
+#         logger.error(f"An error occurred: {e}")
+#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
+# # Endpoint to read embeddings
+# @app.post("/read_embeddings")
+# async def read_embeddings(
+#     request: ReadEmbeddingRequest,
+#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
+# ):
+#     """
+#     Read embeddings from a Hugging Face dataset.
+#     """
+#     try:
+#         df = await huggingface_service.read_dataset(request.dataset_name)
+#         return df
+#     except DatasetNotFoundError as e:
+#         logger.error(f"Dataset not found: {e}")
+#         raise HTTPException(status_code=404, detail=f"Dataset not found: {e}")
+#     except Exception as e:
+#         logger.error(f"An error occurred: {e}")
+#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
+# # 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 by generating embeddings for new data and concatenating it with the existing dataset.
+#     """
+#     try:
+#         # 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}")
+#         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}")
+
+
+# # Endpoint to delete embeddings
+# @app.post("/delete_embeddings")
+# async def delete_embeddings(
+#     request: DeleteEmbeddingRequest,
+#     huggingface_service: HuggingFaceService = Depends(get_huggingface_service),
+# ):
+#     """
+#     Delete embeddings from a Hugging Face dataset.
+#     """
+#     try:
+#         await huggingface_service.delete_dataset(request.dataset_name)
+#         return {
+#             "message": "Embeddings deleted successfully.",
+#             "dataset_name": request.dataset_name,
+#         }
+#     except DatasetPushError as e:
+#         logger.error(f"Failed to delete columns: {e}")
+#         raise HTTPException(status_code=500, detail=f"Failed to delete columns: {e}")
+#     except Exception as e:
+#         logger.error(f"An error occurred: {e}")
+#         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
+
+
 import os
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.responses import JSONResponse, RedirectResponse
 from fastapi.middleware.gzip import GZipMiddleware
 from pydantic import BaseModel
 from typing import List, Dict
+from datasets import Dataset
 from src.api.models.embedding_models import (
     CreateEmbeddingRequest,
     ReadEmbeddingRequest,
@@ -16,8 +266,6 @@ 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
 
@@ -118,21 +366,21 @@ async def create_embedding(
         # Step 1: Query the database
         logger.info("Fetching data from the database...")
         result = await db.fetch(request.query)
-        df = pd.DataFrame(result)
+        dataset = Dataset.from_dict(result)
 
         # Step 2: Generate embeddings
-        df = await embedding_service.create_embeddings(
-            df, request.target_column, request.output_column
+        dataset = await embedding_service.create_embeddings(
+            dataset, request.target_column, request.output_column
         )
 
         # Step 3: Push to Hugging Face Hub
-        await huggingface_service.push_to_hub(df, request.dataset_name)
+        await huggingface_service.push_to_hub(dataset, request.dataset_name)
 
         return JSONResponse(
             content={
                 "message": "Embeddings created and pushed to Hugging Face Hub.",
                 "dataset_name": request.dataset_name,
-                "num_rows": len(df),
+                "num_rows": len(dataset),
             }
         )
     except QueryExecutionError as e:
@@ -159,8 +407,8 @@ async def read_embeddings(
     Read embeddings from a Hugging Face dataset.
     """
     try:
-        df = await huggingface_service.read_dataset(request.dataset_name)
-        return df
+        dataset = await huggingface_service.read_dataset(request.dataset_name)
+        return dataset.to_dict()
     except DatasetNotFoundError as e:
         logger.error(f"Dataset not found: {e}")
         raise HTTPException(status_code=404, detail=f"Dataset not found: {e}")
@@ -170,30 +418,6 @@ 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,
@@ -204,7 +428,7 @@ async def update_embeddings(
     """
     try:
         # Call the update_dataset method to generate embeddings, concatenate, and push the updated dataset
-        updated_df = await huggingface_service.update_dataset(
+        updated_dataset = await huggingface_service.update_dataset(
             request.dataset_name,
             request.updates,
             request.target_column,
@@ -214,7 +438,7 @@ async def update_embeddings(
         return {
             "message": "Embeddings updated successfully.",
             "dataset_name": request.dataset_name,
-            "num_rows": len(updated_df),
+            "num_rows": len(updated_dataset),
         }
     except DatasetPushError as e:
         logger.error(f"Failed to update dataset: {e}")