backend/schema.py

from datetime import datetime
from typing import Dict, List, Any, Optional


def create_submission_record(
    model_name: str,
    hf_space_tag: str,
    model_description: str,
    organization: str,
    model_size: int,  # Changed from str to int
    publication_title: str,
    publication_link: str,

    pretrained: bool,
    pretraining_data: str,
    zero_shot: bool,
    few_shot: bool,
    n_shot: str,
    raw_predictions: List[Dict[str, Any]],
    computed_metrics: Dict[str, Any],
    status: str = "completed",
    approved: bool = False
) -> Dict[str, Any]:
    """Create a standardized submission record for the HuggingFace dataset."""
    
    now = datetime.now().isoformat()
    
    return {
        "config": {
            "model_name": model_name,
            "hf_space_tag": hf_space_tag,
            "model_description": model_description,
            "publication_title": publication_title,
            "publication_link": publication_link,
            "model_size": model_size,
            "pretrained": pretrained,
            "organization": organization,
            "date_submitted": now,
            "date_approved": None,
            "status": status,
            "approved": approved,

            "pretraining_data": pretraining_data,
            "zero_shot": zero_shot,
            "few_shot": few_shot,
            "n_shot": n_shot,
        },
        "raw_predictions": raw_predictions,
        "results": computed_metrics
    }


def get_dataset_schema() -> Dict[str, Any]:
    """Return the HuggingFace dataset schema."""

    return {
        "config": {
            "model_name": "string",
            "hf_space_tag": "string",
            "model_description": "string",
            "publication_title": "string",
            "publication_link": "string",
            "model_size": "int64",  # Changed from "string" to "int64" for numeric values
            "pretraining": "string",
            "organization": "string",
            "date_submitted": "string",
            "date_approved": "string",
            "status": "string",
            "approved": "bool"
        },
        "raw_predictions": "list",
        "results": "dict"
    }