generate_data.py

"""
Generate the OKF Screener dataset from the Hugging Face Hub API.

Fetches the top 50 models by downloads and screens each against the two
EU AI Act Article 53 obligations that apply to ALL general-purpose AI models
— including open-source ones:

  53(1)(c) — Copyright compliance policy
  53(1)(d) — Training data summary

The compliance check uses real model card signals (declared datasets,
README mentions of training data and copyright) rather than just the
model's output license.

Outputs:
  models.csv      — flat tabular format (for the CSV lane)
  models.okf.json — nested JSON-LD with semantic ontology (for the OKF lane)
"""

import json
import re

from huggingface_hub import HfApi, hf_hub_download


# ── License classifications ─────────────────────────────────────────────

PERMISSIVE = {
    "apache-2.0", "mit", "cc-by-4.0", "cc-by-sa-4.0",
    "bsd-2-clause", "bsd-3-clause", "isc", "artistic-2.0",
    "0bsd", "unlicense", "cc0-1.0", "wtfpl", "zlib",
}

RESTRICTIVE = {
    "cc-by-nc-4.0", "cc-by-nc-sa-4.0", "cc-by-nc-nd-4.0",
    "other", "bigscience-bloom-rail-1.0", "bigscience-openrail-m",
    "creativeml-openrail-m",
}


# ── Helpers ──────────────────────────────────────────────────────────────

def parse_param_count(model_id: str) -> tuple[str, float | None]:
    """Extract parameter count from model name (e.g. '7B', '82M')."""
    name = model_id.split("/")[-1]
    match = re.search(r"(\d+\.?\d*)\s*([BbMm])", name)
    if match:
        num = float(match.group(1))
        unit = match.group(2).upper()
        if unit == "M":
            display = f"{int(num)}M" if num == int(num) else f"{num}M"
            return display, num / 1000
        display = f"{int(num)}B" if num == int(num) else f"{num}B"
        return display, num
    return "Unknown", None


def _read_model_card(model_id: str) -> str:
    """Download and return the README.md text, or empty string on failure."""
    try:
        path = hf_hub_download(model_id, "README.md")
        with open(path) as f:
            return f.read()
    except Exception:
        return ""


def get_license_url(spdx: str) -> str:
    urls = {
        "apache-2.0": "https://www.apache.org/licenses/LICENSE-2.0",
        "mit": "https://opensource.org/licenses/MIT",
        "cc-by-4.0": "https://creativecommons.org/licenses/by/4.0/",
        "cc-by-sa-4.0": "https://creativecommons.org/licenses/by-sa/4.0/",
        "cc-by-nc-4.0": "https://creativecommons.org/licenses/by-nc/4.0/",
    }
    return urls.get(spdx, "")


# ── Article 53 compliance checks ─────────────────────────────────────────
#
# These two obligations apply to ALL GPAI providers, including open-source.
# Article 53(2) exempts OSS from 53(1)(a) and (b), but NOT (c) and (d).


def check_copyright_policy(readme: str, license_str: str) -> tuple[str, str]:
    """
    Article 53(1)(c) — Has the provider put in place a copyright policy?

    We check whether the model card mentions copyright-related terms.
    The MODEL LICENSE (apache-2.0, etc.) is about the output weights —
    it tells us nothing about whether the provider respected copyright
    during training.
    """
    readme_lower = readme.lower()

    copyright_keywords = [
        "copyright", "intellectual property", "copyrighted material",
        "opt-out", "opt out", "content removal", "data rights",
        "rights holder", "dmca", "takedown",
    ]
    mentions = [kw for kw in copyright_keywords if kw in readme_lower]

    if len(mentions) >= 2:
        return "Documented", f"Model card references copyright-related terms: {', '.join(mentions[:3])}."
    elif len(mentions) == 1:
        return "Partial", f"Model card has a single mention of '{mentions[0]}' but no detailed copyright compliance policy."
    elif readme:
        return "Undisclosed", "Model card exists but contains no copyright compliance policy or opt-out mechanism."
    else:
        return "No Model Card", "No model card found — impossible to assess copyright compliance."


def check_training_data_summary(readme: str, datasets_declared: list | None) -> tuple[str, str]:
    """
    Article 53(1)(d) — Has the provider published a training data summary?

    We check two signals:
      1. Does the model card metadata declare training datasets?
      2. Does the README mention training data / training corpus?
    """
    readme_lower = readme.lower()

    has_datasets_field = bool(datasets_declared)
    training_keywords = [
        "training data", "training corpus", "training set",
        "pre-training data", "pretraining data", "fine-tuning data",
        "trained on", "fine-tuned on",
    ]
    mentions_training = any(kw in readme_lower for kw in training_keywords)

    if has_datasets_field and mentions_training:
        ds_list = ", ".join(datasets_declared[:4])
        return "Documented", f"Training datasets declared ({ds_list}) and discussed in model card."
    elif has_datasets_field:
        ds_list = ", ".join(datasets_declared[:4])
        return "Partial", f"Datasets listed in metadata ({ds_list}) but no detailed summary in model card."
    elif mentions_training:
        return "Partial", "Model card mentions training data but no structured dataset declarations."
    elif readme:
        return "Undisclosed", "Model card exists but contains no training data information."
    else:
        return "No Model Card", "No model card found — impossible to assess training data transparency."


# ── Main generator ───────────────────────────────────────────────────────

def generate_dataset():
    api = HfApi()
    print("Fetching top 50 models from Hugging Face...")
    models = list(api.list_models(sort="downloads", limit=50))

    csv_lines = [
        "model_id,author,parameters,license,open_weights,"
        "art53c_copyright_policy,art53d_training_data"
    ]

    okf_data = {
        "@context": {
            "schema": "https://schema.org/",
            "eu_ai_act": "https://eur-lex.europa.eu/eli/reg/2024/1689/",
            "model_id": "schema:identifier",
            "author": "schema:creator",
            "architecture": {
                "@id": "schema:applicationCategory",
                "parameters": {
                    "@id": "schema:memoryRequirements",
                    "count": "schema:value",
                    "unit": "schema:unitText",
                },
            },
            "license": {
                "@id": "schema:license",
                "spdx_id": "schema:identifier",
                "url": "schema:url",
                "is_permissive": "schema:isAccessibleForFree",
            },
            "eu_compliance": {
                "@id": "eu_ai_act:Article53",
                "@note": "53(1)(c) and (d) apply to ALL GPAI providers including open-source (53(2) only exempts (a) and (b))",
                "copyright_policy": {
                    "status": "eu_ai_act:complianceStatus",
                    "regulation": "eu_ai_act:Article53_1_c",
                    "evidence": "schema:description",
                },
                "training_data_summary": {
                    "status": "eu_ai_act:complianceStatus",
                    "regulation": "eu_ai_act:Article53_1_d",
                    "evidence": "schema:description",
                },
            },
            "open_weights": "schema:isAccessibleForFree",
        },
        "data": [],
    }

    for i, m in enumerate(models):
        model_id = m.modelId
        author = model_id.split("/")[0] if "/" in model_id else (m.author or "unknown")
        param_display, param_numeric = parse_param_count(model_id)

        # Extract license from tags
        license_str = "unknown"
        for t in (m.tags or []):
            if t.startswith("license:"):
                license_str = t.split(":", 1)[1]
                break

        is_open = license_str not in ("unknown", "proprietary")

        # Read the actual model card for real compliance signals
        print(f"  [{i+1}/50] {model_id}...", end=" ", flush=True)
        readme = _read_model_card(model_id)
        datasets_declared = None
        if m.card_data:
            datasets_declared = getattr(m.card_data, "datasets", None)

        # Article 53(1)(c) — copyright policy
        cr_status, cr_evidence = check_copyright_policy(readme, license_str)

        # Article 53(1)(d) — training data summary
        td_status, td_evidence = check_training_data_summary(readme, datasets_declared)

        print(f"53c={cr_status}, 53d={td_status}")

        # CSV row (flat)
        csv_lines.append(
            f"{model_id},{author},{param_display},{license_str},"
            f"{is_open},{cr_status},{td_status}"
        )

        # OKF record (nested, semantic)
        okf_data["data"].append({
            "model_id": model_id,
            "author": author,
            "architecture": {
                "parameters": {
                    "count": param_numeric,
                    "unit": "billion" if param_numeric and param_numeric >= 1
                            else ("million" if param_numeric else None),
                    "display": param_display,
                    "source": "model_name" if param_display != "Unknown" else "unavailable",
                },
            },
            "license": {
                "spdx_id": license_str,
                "url": get_license_url(license_str),
                "is_permissive": license_str in PERMISSIVE,
            },
            "eu_compliance": {
                "copyright_policy": {
                    "status": cr_status,
                    "regulation": "EU AI Act Article 53(1)(c)",
                    "evidence": cr_evidence,
                },
                "training_data_summary": {
                    "status": td_status,
                    "regulation": "EU AI Act Article 53(1)(d)",
                    "evidence": td_evidence,
                },
            },
            "open_weights": is_open,
        })

    # Write outputs
    with open("models.csv", "w") as f:
        f.write("\n".join(csv_lines))

    with open("models.okf.json", "w") as f:
        json.dump(okf_data, f, indent=2)

    # Print stats
    data = okf_data["data"]
    print(f"\nGenerated models.csv and models.okf.json with {len(data)} records.")

    for field, label in [("copyright_policy", "53(1)(c) Copyright Policy"),
                         ("training_data_summary", "53(1)(d) Training Data")]:
        counts = {}
        for r in data:
            s = r["eu_compliance"][field]["status"]
            counts[s] = counts.get(s, 0) + 1
        parts = [f"{v} {k}" for k, v in sorted(counts.items(), key=lambda x: -x[1])]
        print(f"  {label}: {', '.join(parts)}")


if __name__ == "__main__":
    generate_dataset()