utils/io.py

from huggingface_hub import list_models, model_info
from datetime import datetime
from datasets import Dataset, load_dataset
import pandas as pd
import os
import globals
from typing import List, Tuple


def get_models_providers() -> List[Tuple[str, List[str]]]:
    """Get list of popular text generation models and associated providers from Hugging Face"""
    models = list_models(
        filter="text-generation",
        sort="likes",
        direction=-1,
        limit=globals.NUM_MODELS_RUN,
        expand="inferenceProviderMapping"
    )

    model_providers = [
        (model.id, [p.provider for p in model.inference_provider_mapping])
        for model in models
        if hasattr(model, 'inference_provider_mapping') and model.inference_provider_mapping
    ]
    return model_providers


def initialize_models_providers_file(file_path: str = globals.LOCAL_CONFIG_FILE) -> str:
    """Initialize the models_providers.txt file with popular models and their providers."""
    model_to_providers = get_models_providers()

    with open(file_path, 'w') as f:
        f.write("# Models and Providers Configuration\n")
        f.write("# Format: model_name  provider_name\n")
        f.write(f"# Auto-generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")

        count = 0
        for (model_id, providers) in model_to_providers:
            try:
                for provider in providers:
                    f.write(f"{model_id}  {provider}\n")
                    count += 1
            except Exception as e:
                print(f"Error processing model {model_id}: {e}")
                continue

        print(f"Successfully wrote {count} model-provider combinations to {file_path}")
        return f"Initialized {count} model-provider combinations"


def load_models_providers(file_path: str = globals.LOCAL_CONFIG_FILE) -> List[Tuple[str, str]]:
    """Load models and providers from text file."""
    models_providers = []
    try:
        with open(file_path, 'r') as f:
            for line in f:
                line = line.strip()
                # Skip empty lines and comments
                if line and not line.startswith('#'):
                    parts = line.split()
                    if len(parts) >= 2:
                        model = parts[0]
                        provider = parts[1]
                        models_providers.append((model, provider))
    except Exception as e:
        print(f"Error loading model providers: {str(e)}")
    return models_providers


def save_results() -> None:
    """Persist job results to HuggingFace dataset."""
    try:
        if not globals.job_results:
            print("No results to save")
            return

        records = list(globals.job_results.values())
        df = pd.DataFrame(records)
        dataset = Dataset.from_pandas(df)

        # Push to HuggingFace Hub
        dataset.push_to_hub(
            globals.RESULTS_DATASET_NAME,
            token=os.getenv("HF_TOKEN"),
            private=False
        )
        print(f"Saved {len(records)} results to dataset")

    except Exception as e:
        print(f"Error saving results to dataset: {e}")


def load_results() -> None:
    """Load job results from HuggingFace dataset."""
    try:
        # Try to load existing dataset
        dataset = load_dataset(
            globals.RESULTS_DATASET_NAME,
            split="train",
            token=os.getenv("HF_TOKEN")
        )

        # Convert dataset to job_results dict
        for row in dataset:
            key = globals.get_model_provider_key(row["model"], row["provider"])
            globals.job_results[key] = {
                "model": row["model"],
                "provider": row["provider"],
                "last_run": row["last_run"],
                "status": row["status"],
                "current_score": row["current_score"],
                "previous_score": row["previous_score"],
                "job_id": row["job_id"]
            }

        print(f"Loaded {len(globals.job_results)} results from dataset")

    except Exception as e:
        print(f"No existing dataset found or error loading: {e}")
        print("Starting with empty results")

def style_status(val):
    """Style function for status column."""
    if val == "COMPLETED":
        return 'background-color: green'
    elif val == "ERROR":
        return 'background-color: red'
    elif val == "RUNNING":
        return 'background-color: blue'
    return ''

def get_results_table():
    """Return job results as a styled pandas DataFrame for Gradio DataFrame."""
    if not globals.job_results:
        return pd.DataFrame(columns=["Model", "Provider", "Last Run", "Status", "Current Score", "Previous Score", "Latest Job Id"])

    table_data = []
    for key, info in globals.job_results.items():
        current_score = info.get("current_score", "N/A")
        if current_score is not None and isinstance(current_score, (int, float)):
            current_score = f"{current_score:.4f}"

        previous_score = info.get("previous_score", "N/A")
        if previous_score is not None and isinstance(previous_score, (int, float)):
            previous_score = f"{previous_score:.4f}"

        job_id = info.get("job_id", "N/A")
        # Create a clickable link for the job ID
        if job_id != "N/A":
            job_url = f"https://hf.co/jobs/{globals.NAMESPACE}/{job_id}"
            job_link = f'{job_id}: <a href="{job_url}" target="_blank">📄</a> '
        else:
            job_link = job_id

        # Create relaunch link with data attributes for JavaScript to access
        model = info["model"]
        provider = info["provider"]
        relaunch_link = '🔄 Relaunch'

        table_data.append([
            model,
            provider,
            info["last_run"],
            info["status"],
            current_score,
            previous_score,
            job_link,
            relaunch_link
        ])

    df = pd.DataFrame(table_data, columns=["Model", "Provider", "Last Run", "Status", "Current Score", "Previous Score", "Job Id and Logs", "Actions"])

    # Apply styling to the Status column
    styled_df = df.style.map(style_status, subset=['Status'])
    return styled_df