src/utils.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from datasets import Dataset, load_dataset


def load_image_split(
    dataset_name: str,
    split: str,
    max_samples: int | None = None,
) -> Dataset:
    dataset = load_dataset(dataset_name, split=split)
    if max_samples is not None and max_samples > 0:
        dataset = dataset.select(range(min(max_samples, len(dataset))))
    return dataset


def detect_label_names(dataset: Dataset, label_column: str) -> list[str] | None:
    feature = dataset.features.get(label_column)
    if feature is not None and hasattr(feature, "names"):
        return list(feature.names)
    return None


def label_to_name(label: Any, label_names: list[str] | None) -> str:
    if label_names and isinstance(label, (int, np.integer)) and 0 <= int(label) < len(label_names):
        return label_names[int(label)]
    return str(label)


def embeddings_to_frame(
    embeddings: np.ndarray,
    labels: list[Any],
    label_names: list[str] | None,
    sample_ids: list[str],
    model_name: str,
    dataset_name: str,
    split: str,
) -> pd.DataFrame:
    return pd.DataFrame(
        {
            "sample_id": sample_ids,
            "label": labels,
            "label_name": [label_to_name(label, label_names) for label in labels],
            "dataset_name": dataset_name,
            "split": split,
            "model_name": model_name,
            "embedding": [embedding.astype(float).tolist() for embedding in embeddings],
        }
    )


def save_embeddings(df: pd.DataFrame, output_dir: str | Path) -> Path:
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    parquet_path = output_path / "embeddings.parquet"
    df.to_parquet(parquet_path, index=False)
    return parquet_path


def read_embeddings(path: str | Path) -> pd.DataFrame:
    df = pd.read_parquet(path)
    required = {"sample_id", "label", "embedding"}
    missing = required.difference(df.columns)
    if missing:
        raise ValueError(f"Embedding file is missing required columns: {sorted(missing)}")
    return df


def embedding_matrix(df: pd.DataFrame) -> np.ndarray:
    if df.empty:
        return np.empty((0, 0), dtype=np.float32)
    return np.vstack(df["embedding"].map(np.asarray).to_list()).astype(np.float32)