hf_operations.py

"""
HuggingFace Operations: Upload data, create PRs, validate schemas.
"""
from huggingface_hub import HfApi, login
import pandas as pd
import json
from pathlib import Path
from jsonschema import validate, ValidationError, Draft7Validator


# Load schema once at module level
SCHEMA_PATH = Path(__file__).parent / "eval.schema.json"
with open(SCHEMA_PATH, 'r') as f:
    EVAL_SCHEMA = json.load(f)
    

def validate_json_against_schema(json_data):
    """
    Validate a JSON object against eval.schema.json.
    
    Args:
        json_data: Dict containing the evaluation data
        
    Returns:
        (bool, str): (is_valid, error_message)
    """
    try:
        validate(instance=json_data, schema=EVAL_SCHEMA)
        return True, "Schema validation passed"
    except ValidationError as e:
        # Extract the most relevant error message
        error_path = " → ".join(str(p) for p in e.path) if e.path else "root"
        return False, f"❌ Schema validation failed at '{error_path}': {e.message}"
    except Exception as e:
        return False, f"❌ Validation error: {str(e)}"


def upload_to_hf_dataset(parquet_file, split_name, repo_id):
    """
    Upload a parquet file as a new split to the HF dataset.
    
    Args:
        parquet_file: Path to parquet file
        split_name: Name of the split (leaderboard name)
        repo_id: HuggingFace dataset repository ID
    """
    # TODO: Implement upload logic
    pass


def check_hf_authentication():
    """
    Check if user is authenticated with HuggingFace.
    
    Returns:
        (bool, str): (is_authenticated, username or error_message)
    """
    try:
        api = HfApi()
        user_info = api.whoami()
        return True, user_info['name']
    except Exception as e:
        return False, "Not authenticated. Run: huggingface-cli login"


def check_duplicate_pr_exists(leaderboard_name, repo_id):
    """
    Check if a PR already exists for this leaderboard.
    
    Args:
        leaderboard_name: Name of the leaderboard
        repo_id: HuggingFace dataset repository ID
        
    Returns:
        (bool, str or None): (exists, pr_url if exists)
    """
    try:
        api = HfApi()
        discussions = api.get_repo_discussions(repo_id=repo_id, repo_type="dataset")
        
        # Check for open PRs with matching title
        pr_title_pattern = f"add new leaderboard: {leaderboard_name.lower()}"
        for discussion in discussions:
            if discussion.is_pull_request and discussion.status == "open":
                if pr_title_pattern in discussion.title.lower():
                    pr_url = f"https://huggingface.co/datasets/{repo_id}/discussions/{discussion.num}"
                    return True, pr_url
        
        return False, None
    except Exception as e:
        # If we can't check, assume no duplicate (fail open)
        print(f"Warning: Could not check for duplicate PRs: {e}")
        return False, None


def create_pr_for_new_leaderboard(leaderboard_name, parquet_file, repo_id):
    """
    Create a pull request to add a new leaderboard split.
    
    Args:
        leaderboard_name: Name of the new leaderboard
        parquet_file: Path to parquet file
        repo_id: HuggingFace dataset repository ID
        
    Returns:
        (success, pr_url or error_message)
    """
    # 1. Check authentication
    is_auth, auth_result = check_hf_authentication()
    if not is_auth:
        return False, f"❌ {auth_result}"
    
    # 2. Check for duplicate PR
    has_duplicate, duplicate_url = check_duplicate_pr_exists(leaderboard_name, repo_id)
    if has_duplicate:
        return False, f"⚠️ PR already exists: {duplicate_url}"
    
    # 3. Validate parquet file exists and has data
    parquet_path = Path(parquet_file)
    if not parquet_path.exists():
        return False, "❌ Parquet file not found"
    
    df = pd.read_parquet(parquet_file)
    if len(df) == 0:
        return False, "❌ Parquet file is empty"
    
    # 4. Create PR
    try:
        api = HfApi()
        
        # Upload the parquet file to the branch
        commit_message = f"Add new leaderboard: {leaderboard_name}"
        
        # Upload file and create PR
        commit_info = api.upload_file(
            path_or_fileobj=parquet_file,
            path_in_repo=f"data/{leaderboard_name}.parquet",
            repo_id=repo_id,
            repo_type="dataset",
            commit_message=commit_message,
            create_pr=True,
        )
        
        # Extract PR URL from commit info
        pr_url = commit_info.pr_url if hasattr(commit_info, 'pr_url') else f"https://huggingface.co/datasets/{repo_id}/discussions"
        
        return True, f"PR created ({len(df)} rows): {pr_url}"
        
    except Exception as e:
        return False, f"❌ Failed to create PR: {str(e)}"


def validate_schema(parquet_file):
    """
    Validate that a parquet file matches the expected schema.
    
    Args:
        parquet_file: Path to parquet file to validate
        
    Returns:
        (bool, str): (is_valid, error_message)
    """
    try:
        df = pd.read_parquet(parquet_file)
        
        # Required columns
        required_cols = [
            '_leaderboard', '_developer', '_model', '_uuid',
            'schema_version', 'evaluation_id', 'retrieved_timestamp',
            'source_data', 'evaluation_source_name', 'evaluation_source_type',
            'source_organization_name', 'evaluator_relationship',
            'model_name', 'model_id', 'model_developer',
            'evaluation_results'
        ]
        
        missing = [col for col in required_cols if col not in df.columns]
        if missing:
            return False, f"Missing required columns: {', '.join(missing)}"
        
        # Check data types (all should be strings)
        for col in df.columns:
            if df[col].dtype not in ['object', 'string']:
                return False, f"Column '{col}' has wrong type: {df[col].dtype} (expected string)"
        
        return True, "Schema validation passed"
    
    except Exception as e:
        return False, f"Validation error: {str(e)}"


def export_to_json(parquet_file, output_dir):
    """
    Export parquet data back to JSON files.
    Uses the parquet_to_folder function from json_to_parquet.py
    
    Args:
        parquet_file: Path to parquet file
        output_dir: Directory to write JSON files to
    """
    from json_to_parquet import parquet_to_folder
    parquet_to_folder(parquet_file, output_dir)