scripts/release_check.py

from __future__ import annotations

import json
import subprocess
from collections import Counter
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
MAX_FILE_SIZE_BYTES = 5_000_000  # 5 MB

REQUIRED = [
    'README.md',
    'app.py',
    'requirements.txt',
    'research_config.json',
    'featurelens/runtime.py',
    'featurelens/sae.py',
    'featurelens/interventions.py',
    'featurelens/metrics.py',
    'featurelens/stats.py',
    'experiments/run_all.py',
    'experiments/run_causal.py',
    'experiments/run_feature_sets.py',
    'data/prompts.jsonl',
    'data/causal_tasks.jsonl',
    'docs/VALIDATION.md',
    'scripts/ui_smoke.py',
]


def load_jsonl(path: Path) -> list[dict]:
    return [
        json.loads(line)
        for line in path.read_text(encoding='utf-8').splitlines()
        if line.strip()
    ]


def repository_candidates() -> list[Path]:
    """Return tracked files plus untracked files that are not ignored by Git."""
    try:
        result = subprocess.run(
            ['git', 'ls-files', '--cached', '--others', '--exclude-standard'],
            cwd=ROOT,
            capture_output=True,
            text=True,
            check=True,
        )
    except FileNotFoundError as exc:
        raise SystemExit('Git is required to run the FeatureLens release check.') from exc
    except subprocess.CalledProcessError as exc:
        raise SystemExit(
            f'Could not inspect repository files with Git: {exc.stderr.strip()}'
        ) from exc

    paths: list[Path] = []
    for relative_path in result.stdout.splitlines():
        relative_path = relative_path.strip()
        if not relative_path:
            continue
        path = ROOT / relative_path
        if path.is_file():
            paths.append(path)
    return paths


def check_required_files() -> None:
    missing = [path for path in REQUIRED if not (ROOT / path).exists()]
    if missing:
        raise SystemExit(f'Missing required files: {missing}')


def check_config(config: dict) -> None:
    expected = {
        'layers': [4, 14, 26],
        'model_id': 'Qwen/Qwen3-1.7B-Base',
        'sae_width': 32768,
        'dose_response_multipliers': [0.0, 0.5, 1.0, 1.5, 2.0, 3.0],
        'feature_set_sizes': [1, 3, 5],
        'live_random_controls': 8,
        'offline_random_controls_default': 8,
        'concept_contrast_prompts_per_concept': 4,
        'interaction_feature_limit': 5,
        'live_geometry_feature_limit': 8,
        'concept_contrast_pooling': 'max activation across non-padding prompt tokens',
    }
    for key, value in expected.items():
        if config.get(key) != value:
            raise SystemExit(f'Unexpected {key}: {config.get(key)!r}. Expected {value!r}.')

    required_live_v04 = {
        'batch_context_null_reference',
        'random_control_ensemble',
        'individual_vs_joint_interaction_decomposition',
        'promptwide_paraphrase_robustness',
        'controlled_concept_contrast_scan',
        'copy_tables_with_headers',
    }
    actual_live_v04 = set(config.get('live_features_v0_4', []))
    if actual_live_v04 != required_live_v04:
        raise SystemExit(
            'research_config.json live_features_v0_4 mismatch: '
            f'{sorted(actual_live_v04)}'
        )

    required_live_v05 = {
        'wide_centered_responsive_layout',
        'copy_feedback',
        'dynamic_height_reflow_observer',
        'promptwide_concept_contrast_scan',
        'feature_token_activation_trace',
        'contrastive_continuation_preference_test',
        'feature_decoder_geometry',
    }
    actual_live_v05 = set(config.get('live_features_v0_5', []))
    if actual_live_v05 != required_live_v05:
        raise SystemExit(
            'research_config.json live_features_v0_5 mismatch: '
            f'{sorted(actual_live_v05)}'
        )

    required_live_v06 = {
        'start_here_plain_language_onboarding',
        'persistent_workbench_context_banner',
        'explicit_per_experiment_feature_selectors',
        'plot_fullscreen_and_export_controls',
        'consistent_heading_and_table_typography',
        'concept_guided_candidate_feature_discovery',
        'completion_cue_sensitivity_scan',
    }
    actual_live_v06 = set(config.get('live_features_v0_6', []))
    if actual_live_v06 != required_live_v06:
        raise SystemExit(
            'research_config.json live_features_v0_6 mismatch: '
            f'{sorted(actual_live_v06)}'
        )

    required_live_v07 = {
        'cleaned_nonaccordion_experiment_layout',
        'focused_fullscreen_modal_for_tables_and_plots',
        'descriptive_plot_export_filenames',
        'german_language_control_concept',
        'balanced_candidate_ranking_and_current_prompt_compatibility',
        'click_to_select_candidate_rows',
        'completion_cue_context_matrix',
    }
    actual_live_v07 = set(config.get('live_features_v0_7', []))
    if actual_live_v07 != required_live_v07:
        raise SystemExit(
            'research_config.json live_features_v0_7 mismatch: '
            f'{sorted(actual_live_v07)}'
        )
    if 'german_language' not in config.get('concepts', []) or 'french_language' in config.get('concepts', []):
        raise SystemExit('research_config.json must use german_language and must not contain french_language.')


def check_datasets(config: dict) -> tuple[list[dict], list[dict]]:
    prompts = load_jsonl(ROOT / 'data' / 'prompts.jsonl')
    causal = load_jsonl(ROOT / 'data' / 'causal_tasks.jsonl')

    if len(prompts) != config.get('discovery_prompts'):
        raise SystemExit(
            f'Discovery prompt count mismatch: found {len(prompts)}, '
            f'expected {config.get("discovery_prompts")}.'
        )
    if len(causal) != config.get('causal_tasks'):
        raise SystemExit(
            f'Causal task count mismatch: found {len(causal)}, '
            f'expected {config.get("causal_tasks")}.'
        )

    concept_counts = Counter(row['concept'] for row in prompts)
    if set(concept_counts) != set(config.get('concepts', [])):
        raise SystemExit('Discovery dataset concepts do not match research_config.json.')
    if len(set(concept_counts.values())) != 1:
        raise SystemExit(f'Discovery concepts are not balanced: {dict(concept_counts)}')

    pair_counts = Counter(row['pair_id'] for row in prompts)
    if set(pair_counts.values()) != {2}:
        raise SystemExit('Every discovery paraphrase pair must contain exactly two prompts.')

    return prompts, causal


def check_oversized_files() -> None:
    oversized: list[str] = []
    for path in repository_candidates():
        size_bytes = path.stat().st_size
        if size_bytes > MAX_FILE_SIZE_BYTES:
            relative = path.relative_to(ROOT)
            oversized.append(f'{relative} ({size_bytes / 1_000_000:.1f} MB)')

    if oversized:
        formatted = '\n  - '.join(oversized)
        raise SystemExit(
            'Repository contains unexpectedly large tracked/unignored candidates:\n'
            f'  - {formatted}\n\n'
            'If a file is a legitimate local artifact, add it to .gitignore. Model weights, '
            'SAE checkpoints, activation dumps, virtual environments, and caches should not be committed.'
        )


def check_readme() -> None:
    readme = (ROOT / 'README.md').read_text(encoding='utf-8')
    required_strings = [
        'sdk: gradio',
        'sdk_version: "6.24.0"',
        'Qwen/Qwen3-1.7B-Base',
        'full-continuation',
        'feature-set',
        'paraphrase',
        'random-control ensemble',
        'batched zero-edit',
        'concept contrast',
        'non-additivity',
        'contrastive',
        'decoder geometry',
        'token activation',
        'concept-guided candidate',
        'completion-cue',
        'cue × context',
        'balanced selectivity',
        'german',
        'start here',
    ]
    missing = [value for value in required_strings if value.lower() not in readme.lower()]
    if missing:
        raise SystemExit(f'README.md is missing required v0.7 content: {missing}')


def check_pyproject() -> None:
    text = (ROOT / 'pyproject.toml').read_text(encoding='utf-8')
    if 'version = "0.7.0"' not in text:
        raise SystemExit('pyproject.toml must declare version 0.7.0.')


def main() -> None:
    check_required_files()
    config = json.loads((ROOT / 'research_config.json').read_text(encoding='utf-8'))
    check_config(config)
    prompts, causal = check_datasets(config)
    check_oversized_files()
    check_readme()
    check_pyproject()

    print('FeatureLens release check: PASS')
    print(f'  discovery prompts: {len(prompts)}')
    print(f'  causal tasks: {len(causal)}')
    print(f'  layers: {config["layers"]}')
    print(f'  feature-set sizes: {config["feature_set_sizes"]}')
    print(f'  random controls: {config["live_random_controls"]}')
    print('  release: v0.7.0')


if __name__ == '__main__':
    main()