Spaces:
Running
Running
| from datasets import Dataset, load_dataset, concatenate_datasets | |
| from huggingface_hub import HfApi, HfFolder | |
| import logging | |
| import os | |
| from typing import Optional, Dict, List | |
| 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, dataset: Dataset, dataset_name: str) -> None: | |
| """Push the dataset to Hugging Face Hub.""" | |
| try: | |
| logger.info(f"Creating Hugging Face Dataset: {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[Dataset]: | |
| """Read a dataset from Hugging Face Hub.""" | |
| try: | |
| logger.info(f"Loading dataset from Hugging Face Hub: {dataset_name}...") | |
| dataset = load_dataset( | |
| dataset_name, | |
| keep_in_memory=True, | |
| download_mode="force_redownload", | |
| ) | |
| return dataset["train"] | |
| 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[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_dataset = await self.read_dataset(dataset_name) | |
| # 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_dataset = await embedding_service.create_embeddings( | |
| new_dataset, target_column, output_column | |
| ) | |
| # Step 4: Concatenate the existing Dataset with the new Dataset | |
| logger.info("Concatenating existing dataset with new data...") | |
| 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_dataset, dataset_name) | |
| return updated_dataset | |
| 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}") | |