Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Download a structured_analysis parquet and build the local grading template. | |
| Default target is the tiny development run:: | |
| uv run python hf-spaces/grading-answers/ingest_structured_analysis.py | |
| That pulls the latest ``new-llm-judge/tiny-example/en_us/*.parquet`` from | |
| ``TransLegal/structured_analysis`` (branch ``dev``) and writes:: | |
| hf-spaces/grading-answers/data/en-us/grading_template.parquet | |
| Pass ``--input`` to skip Hugging Face and build from a local parquet instead. | |
| Use ``--poll`` while a Dagster run is still publishing. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import logging | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import pandas as pd | |
| from huggingface_hub import hf_hub_download | |
| _REPO_ROOT = Path(__file__).resolve().parents[2] | |
| if str(_REPO_ROOT) not in sys.path: | |
| sys.path.insert(0, str(_REPO_ROOT)) | |
| # Sibling module (same directory). | |
| from build_grading_template import ( # noqa: E402 | |
| build_grading_template, | |
| build_template_metadata, | |
| log_concept_balance, | |
| log_shortfall, | |
| write_template_metadata, | |
| ) | |
| from config import get_settings # noqa: E402 | |
| from config.constants import REPOS # noqa: E402 | |
| from utils.huggingface import get_hf_api_wrapper # noqa: E402 | |
| logger = logging.getLogger(__name__) | |
| _SPACE_ROOT = Path(__file__).resolve().parent | |
| _DEFAULT_OUTPUT = _SPACE_ROOT / "data" / "en-us" / "grading_template.parquet" | |
| _STRUCTURED_ANALYSIS_REPO = REPOS["structured_answers"] | |
| def _hyphen_to_underscore(jurisdiction: str) -> str: | |
| return jurisdiction.replace("-", "_") | |
| def _list_run_parquets( | |
| *, | |
| branch: str, | |
| experiment: str, | |
| variant: str, | |
| jurisdiction: str, | |
| ) -> list[str]: | |
| """Return matching parquet paths under experiment/variant/jurisdiction/.""" | |
| prefix = f"{experiment}/{variant}/{_hyphen_to_underscore(jurisdiction)}/" | |
| hf = get_hf_api_wrapper() | |
| files = hf.api.list_repo_files( | |
| repo_id=f"{hf.config.organization}/{_STRUCTURED_ANALYSIS_REPO}", | |
| repo_type="dataset", | |
| revision=branch, | |
| ) | |
| return sorted(f for f in files if f.startswith(prefix) and f.endswith(".parquet")) | |
| def resolve_hf_parquet( | |
| *, | |
| branch: str, | |
| experiment: str, | |
| variant: str, | |
| jurisdiction: str, | |
| path: str | None, | |
| ) -> str: | |
| """Return the repo-relative parquet path to download.""" | |
| if path: | |
| return path | |
| matches = _list_run_parquets( | |
| branch=branch, | |
| experiment=experiment, | |
| variant=variant, | |
| jurisdiction=jurisdiction, | |
| ) | |
| if not matches: | |
| raise FileNotFoundError( | |
| f"No structured_analysis parquet yet under " | |
| f"{experiment}/{variant}/{_hyphen_to_underscore(jurisdiction)}/ " | |
| f"on branch {branch!r}. Wait for the pipeline publish, or pass --input." | |
| ) | |
| chosen = matches[-1] | |
| if len(matches) > 1: | |
| logger.info("Found %d parquet(s); using latest: %s", len(matches), chosen) | |
| return chosen | |
| def download_hf_parquet( | |
| *, | |
| repo_path: str, | |
| branch: str, | |
| local_dir: Path, | |
| ) -> Path: | |
| settings = get_settings() | |
| org = settings.huggingface.organization | |
| token = settings.huggingface.token | |
| local_dir.mkdir(parents=True, exist_ok=True) | |
| downloaded = hf_hub_download( | |
| repo_id=f"{org}/{_STRUCTURED_ANALYSIS_REPO}", | |
| filename=repo_path, | |
| repo_type="dataset", | |
| revision=branch, | |
| token=token, | |
| local_dir=str(local_dir), | |
| ) | |
| return Path(downloaded) | |
| def poll_until_available( | |
| *, | |
| branch: str, | |
| experiment: str, | |
| variant: str, | |
| jurisdiction: str, | |
| path: str | None, | |
| interval_s: float, | |
| timeout_s: float, | |
| ) -> str: | |
| """Block until a matching parquet appears on HF (or timeout).""" | |
| deadline = time.monotonic() + timeout_s | |
| last_error: Exception | None = None | |
| while time.monotonic() < deadline: | |
| try: | |
| return resolve_hf_parquet( | |
| branch=branch, | |
| experiment=experiment, | |
| variant=variant, | |
| jurisdiction=jurisdiction, | |
| path=path, | |
| ) | |
| except FileNotFoundError as exc: | |
| last_error = exc | |
| remaining = max(0, int(deadline - time.monotonic())) | |
| logger.info( | |
| "Waiting for HF publish (%ds left)… %s", | |
| remaining, | |
| exc, | |
| ) | |
| time.sleep(interval_s) | |
| raise TimeoutError( | |
| f"Timed out after {timeout_s:.0f}s waiting for structured_analysis on " | |
| f"branch {branch!r}." | |
| ) from last_error | |
| def build_and_write( | |
| *, | |
| input_parquet: Path, | |
| output_parquet: Path, | |
| all_fields_of_law: bool, | |
| max_terms: int | None, | |
| include_failed: bool, | |
| terms_per_field: int | None, | |
| jurisdiction: str, | |
| ) -> pd.DataFrame: | |
| logger.info("Loading %s", input_parquet) | |
| structured = pd.read_parquet(input_parquet) | |
| build = build_grading_template( | |
| structured, | |
| only_ok=not include_failed, | |
| balanced=not all_fields_of_law, | |
| terms_per_field=terms_per_field, | |
| max_terms=max_terms, | |
| ) | |
| template = build.template | |
| output_parquet.parent.mkdir(parents=True, exist_ok=True) | |
| template.to_parquet(output_parquet, index=False) | |
| metadata = build_template_metadata( | |
| input_parquet=input_parquet, | |
| output_parquet=output_parquet, | |
| build=build, | |
| structured_input=structured, | |
| only_ok=not include_failed, | |
| jurisdiction=jurisdiction, | |
| max_terms=max_terms, | |
| ) | |
| meta_path = write_template_metadata(metadata, output_parquet) | |
| log_concept_balance(metadata["selection"]["concept_balance"]) | |
| if build.assignment is not None: | |
| log_shortfall(metadata["selection"]["shortfall"]) | |
| n_results = ( | |
| template[["term", "field_of_law"]].drop_duplicates().shape[0] | |
| if not template.empty | |
| else 0 | |
| ) | |
| logger.info( | |
| "Wrote %d grading rows for %d results → %s", | |
| len(template), | |
| n_results, | |
| output_parquet, | |
| ) | |
| logger.info("Wrote metadata → %s", meta_path) | |
| if not template.empty: | |
| logger.info("Concepts: %s", sorted(template["concept"].unique())) | |
| logger.info( | |
| "Questions per concept: %s", | |
| template.groupby("concept")["question_id"].nunique().to_dict(), | |
| ) | |
| return template | |
| def _parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser( | |
| description=( | |
| "Ingest structured_analysis into the local grading app " | |
| "(download from HF and/or build grading_template.parquet)." | |
| ) | |
| ) | |
| parser.add_argument( | |
| "--input", | |
| type=Path, | |
| default=None, | |
| help="Local structured_analysis parquet (skips Hugging Face download).", | |
| ) | |
| parser.add_argument( | |
| "--output", | |
| type=Path, | |
| default=_DEFAULT_OUTPUT, | |
| help=f"grading_template.parquet path (default: {_DEFAULT_OUTPUT}).", | |
| ) | |
| parser.add_argument( | |
| "--branch", | |
| default="dev", | |
| help="HF dataset branch (default: dev).", | |
| ) | |
| parser.add_argument( | |
| "--experiment", | |
| default="new-llm-judge", | |
| help="Experiment folder on HF (default: new-llm-judge).", | |
| ) | |
| parser.add_argument( | |
| "--variant", | |
| default="tiny-example", | |
| help="Variant folder on HF (default: tiny-example).", | |
| ) | |
| parser.add_argument( | |
| "--jurisdiction", | |
| default="en-us", | |
| help="Jurisdiction code with hyphen (default: en-us).", | |
| ) | |
| parser.add_argument( | |
| "--hf-path", | |
| default=None, | |
| help=( | |
| "Exact repo-relative parquet path. Default: latest under " | |
| "experiment/variant/jurisdiction/." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--download-dir", | |
| type=Path, | |
| default=_SPACE_ROOT / "data" / "_downloads", | |
| help="Where to cache the downloaded HF parquet.", | |
| ) | |
| parser.add_argument( | |
| "--poll", | |
| action="store_true", | |
| help="Keep checking HF until a matching parquet appears.", | |
| ) | |
| parser.add_argument( | |
| "--poll-interval", | |
| type=float, | |
| default=30.0, | |
| help="Seconds between HF checks when --poll is set (default: 30).", | |
| ) | |
| parser.add_argument( | |
| "--poll-timeout", | |
| type=float, | |
| default=3600.0, | |
| help="Max seconds to wait with --poll (default: 3600).", | |
| ) | |
| parser.add_argument( | |
| "--all-fields-of-law", | |
| action="store_true", | |
| help="Keep every (term, field_of_law) row (default: one field per term).", | |
| ) | |
| parser.add_argument( | |
| "--max-terms", | |
| type=int, | |
| default=None, | |
| help=( | |
| "Cap the number of terms graded, taken evenly across concept " | |
| "classes (default: every term)." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--include-failed", | |
| action="store_true", | |
| help="Include generation_status != ok rows.", | |
| ) | |
| parser.add_argument( | |
| "--terms-per-field", | |
| type=int, | |
| default=None, | |
| help="Minimum terms per field of law (default: terms // fields).", | |
| ) | |
| args = parser.parse_args() | |
| if args.max_terms is not None and args.max_terms < 0: | |
| parser.error("--max-terms must not be negative") | |
| return args | |
| def main() -> None: | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s %(levelname)s %(name)s: %(message)s", | |
| ) | |
| args = _parse_args() | |
| if args.input is not None: | |
| input_parquet = args.input.expanduser().resolve() | |
| if not input_parquet.is_file(): | |
| raise FileNotFoundError(f"Input parquet not found: {input_parquet}") | |
| else: | |
| if args.poll: | |
| repo_path = poll_until_available( | |
| branch=args.branch, | |
| experiment=args.experiment, | |
| variant=args.variant, | |
| jurisdiction=args.jurisdiction, | |
| path=args.hf_path, | |
| interval_s=args.poll_interval, | |
| timeout_s=args.poll_timeout, | |
| ) | |
| else: | |
| repo_path = resolve_hf_parquet( | |
| branch=args.branch, | |
| experiment=args.experiment, | |
| variant=args.variant, | |
| jurisdiction=args.jurisdiction, | |
| path=args.hf_path, | |
| ) | |
| logger.info( | |
| "Downloading hf://datasets/.../%s@%s → %s", | |
| repo_path, | |
| args.branch, | |
| args.download_dir, | |
| ) | |
| input_parquet = download_hf_parquet( | |
| repo_path=repo_path, | |
| branch=args.branch, | |
| local_dir=args.download_dir, | |
| ) | |
| build_and_write( | |
| input_parquet=input_parquet, | |
| output_parquet=args.output.expanduser().resolve(), | |
| all_fields_of_law=args.all_fields_of_law, | |
| max_terms=args.max_terms, | |
| include_failed=args.include_failed, | |
| terms_per_field=args.terms_per_field, | |
| jurisdiction=args.jurisdiction, | |
| ) | |
| logger.info( | |
| "App-ready. From hf-spaces/grading-answers run: " | |
| "uv run --with streamlit --with streamlit-shortcuts --with pandas " | |
| "--with pyarrow --with huggingface-hub --with pycountry --with babel " | |
| "--with python-dotenv streamlit run src/streamlit_app.py" | |
| ) | |
| if __name__ == "__main__": | |
| main() | |