backend/dataset_storage.py

from datasets import Dataset, load_dataset
from typing import Dict, Any, List
from config.settings import RESULTS_DATASET, HF_TOKEN


def save_submission_to_dataset(submission_record: Dict[str, Any]) -> None:
    """Save a submission record to the HuggingFace dataset."""
    
    # Load existing dataset or create new one
    try:
        dataset = load_dataset(RESULTS_DATASET, token=HF_TOKEN, split="test")
        existing_data = list(dataset)
    except:
        # Dataset doesn't exist or is empty, start fresh
        existing_data = []
    
    # Add new submission
    existing_data.append(submission_record)
    
    # Create new dataset
    new_dataset = Dataset.from_list(existing_data)
    
    # Push to HuggingFace
    new_dataset.push_to_hub(
        RESULTS_DATASET, 
        token=HF_TOKEN,
        split="test"
    )


def get_all_submissions(approved_only: bool = False) -> List[Dict[str, Any]]:
    """Get all submissions from the dataset."""
    
    dataset = load_dataset(RESULTS_DATASET, token=HF_TOKEN, split="test")
    submissions = list(dataset)
    
    if approved_only:
        submissions = [s for s in submissions if s["config"]["approved"]]
    
    return submissions


def update_submission_approval(model_name: str, hf_space_tag: str, approved: bool) -> None:
    """Update the approval status of a submission."""
    
    # Load dataset
    dataset = load_dataset(RESULTS_DATASET, token=HF_TOKEN, split="test")
    data = list(dataset)
    
    # Find and update the submission
    for submission in data:
        config = submission["config"]
        if config["model_name"] == model_name and config["hf_space_tag"] == hf_space_tag:
            config["approved"] = approved
            if approved:
                from datetime import datetime
                config["date_approved"] = datetime.now().isoformat()
            break
    
    # Save updated dataset
    updated_dataset = Dataset.from_list(data)
    updated_dataset.push_to_hub(
        RESULTS_DATASET,
        token=HF_TOKEN, 
        split="test"
    )