captchaforge 0.2.40

#!/usr/bin/env python3
"""Knowledge-distill a captchaforge VLM into a small CPU-runnable
student model.

The fine-tuned 7B teacher (output of train_lora.py) is too heavy
for production inference on CPU-only edge nodes. This script
trains a small (target: 0.5B params) student to mimic the
teacher's outputs on a large body of synthetic + real captcha
challenges.

Approach: standard offline distillation.

  1. Run the teacher on every sample in --teacher-dataset, capture
     its (prompt, answer) outputs to a distillation manifest.
  2. Fine-tune a smaller base model (default: Qwen2-VL-2B) on
     that manifest with the same train_lora.py recipe.
  3. Save the merged student weights ready for promote_to_ollama.

Result: a model that runs in <500MB on disk + <100ms per
inference on a modern CPU. Use as the captchaforge default for
the 99% of captchas the heavier teacher isn't needed for; fall
back to the teacher only when student confidence is below
threshold.
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

LOG = logging.getLogger("captchaforge.distill")


def parse_args(argv: list[str]) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="distill.py",
        description="Distill a captchaforge teacher VLM into a small student.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument(
        "--teacher",
        type=Path,
        required=True,
        help="Path to the teacher model dir (output of train_lora.py).",
    )
    p.add_argument(
        "--teacher-dataset",
        type=Path,
        required=True,
        help="Captchaforge VlmDataset to run the teacher against.",
    )
    p.add_argument(
        "--student-base",
        default="Qwen/Qwen2-VL-2B-Instruct",
        help="HF base model to use as the student.",
    )
    p.add_argument(
        "--out",
        type=Path,
        required=True,
        help="Output dir for the distilled student weights.",
    )
    p.add_argument(
        "--distill-manifest",
        type=Path,
        default=None,
        help="Where to write the teacher-labelled manifest. "
        "Default: <out>/distill-manifest.jsonl",
    )
    p.add_argument(
        "--epochs",
        type=int,
        default=3,
        help="Student fine-tune epochs.",
    )
    p.add_argument(
        "--batch-size",
        type=int,
        default=8,
        help="Smaller model = bigger batch.",
    )
    p.add_argument(
        "--max-samples",
        type=int,
        default=None,
        help="Cap the teacher-labelled corpus.",
    )
    p.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the planned operations; skip actual model load + work.",
    )
    p.add_argument(
        "--log-level",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
    )
    return p.parse_args(argv)


def teacher_label_manifest(
    teacher_dir: Path,
    teacher_dataset: Path,
    out_manifest: Path,
    max_samples: int | None,
) -> int:
    """Run the teacher across `teacher_dataset` and write a new
    manifest with the teacher's outputs as the answer field.
    Returns the count of samples written.
    """
    LOG.info("teacher-labelling %s with model at %s", teacher_dataset, teacher_dir)

    # Imports deferred until we're actually about to run.
    import torch  # noqa: WPS433
    from PIL import Image  # noqa: WPS433
    from transformers import (  # noqa: WPS433
        AutoProcessor,
        Qwen2VLForConditionalGeneration,
    )

    processor = AutoProcessor.from_pretrained(str(teacher_dir))
    model = Qwen2VLForConditionalGeneration.from_pretrained(
        str(teacher_dir),
        torch_dtype=torch.bfloat16,
        device_map="auto",
    )

    src_manifest = teacher_dataset / "manifest.jsonl"
    out_manifest.parent.mkdir(parents=True, exist_ok=True)
    written = 0
    with src_manifest.open("r", encoding="utf-8") as src, out_manifest.open(
        "w", encoding="utf-8"
    ) as dst:
        for line_no, raw in enumerate(src, start=1):
            raw = raw.strip()
            if not raw:
                continue
            sample = json.loads(raw)
            image_path = teacher_dataset / sample["image_path"]
            if not image_path.exists():
                LOG.warning("image missing: %s", image_path)
                continue
            image = Image.open(image_path).convert("RGB")
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image},
                        {"type": "text", "text": sample["prompt"]},
                    ],
                }
            ]
            text = processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
            inputs = processor(text=[text], images=[image], return_tensors="pt").to(
                model.device
            )
            with torch.no_grad():
                generated = model.generate(
                    **inputs,
                    max_new_tokens=256,
                    do_sample=False,  # deterministic teacher labels
                )
            new_tokens = generated[0][inputs.input_ids.shape[1]:]
            answer = processor.decode(new_tokens, skip_special_tokens=True).strip()
            sample["answer"] = answer
            sample["teacher_distilled"] = True
            dst.write(json.dumps(sample))
            dst.write("\n")
            written += 1
            if written % 100 == 0:
                LOG.info("teacher-labelled %d samples", written)
            if max_samples is not None and written >= max_samples:
                break
    LOG.info("teacher-labelling done: %d samples", written)
    return written


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s %(levelname)s %(name)s %(message)s",
    )

    distill_manifest = (
        args.distill_manifest
        if args.distill_manifest is not None
        else args.out / "distill-manifest.jsonl"
    )

    if args.dry_run:
        LOG.info("dry-run plan:")
        LOG.info("  1. teacher-label %s with %s", args.teacher_dataset, args.teacher)
        LOG.info("     -> %s", distill_manifest)
        LOG.info(
            "  2. fine-tune %s on the teacher-labelled manifest",
            args.student_base,
        )
        LOG.info("     -> %s", args.out)
        LOG.info("  3. save merged weights")
        return 0

    args.out.mkdir(parents=True, exist_ok=True)

    LOG.info("step 1/2, teacher-labelling")
    written = teacher_label_manifest(
        args.teacher,
        args.teacher_dataset,
        distill_manifest,
        args.max_samples,
    )
    if written == 0:
        LOG.error("no samples teacher-labelled; aborting")
        return 2

    # Step 2: invoke train_lora.py against the teacher-labelled
    # manifest. We construct a synthetic VlmDataset directory with
    # the manifest + a symlink to the teacher_dataset's images/
    # so train_lora.py's image-resolution path Just Works.
    LOG.info("step 2/2, fine-tuning student on distill manifest")
    student_dataset = args.out / "distill-dataset"
    (student_dataset).mkdir(parents=True, exist_ok=True)
    student_manifest = student_dataset / "manifest.jsonl"
    if student_manifest.exists() or student_manifest.is_symlink():
        student_manifest.unlink()
    student_manifest.symlink_to(distill_manifest.resolve())

    images_link = student_dataset / "images"
    if images_link.exists() or images_link.is_symlink():
        images_link.unlink()
    images_link.symlink_to((args.teacher_dataset / "images").resolve())

    # Hand off to train_lora.py via subprocess so the trainer's
    # full CLI surface (incl. the args we don't expose here) stays
    # available. This keeps the two scripts source-of-truth-singular.
    import subprocess  # noqa: WPS433

    cmd = [
        sys.executable,
        str(Path(__file__).parent / "train_lora.py"),
        "--base-model",
        args.student_base,
        "--dataset",
        str(student_dataset),
        "--out",
        str(args.out / "student"),
        "--epochs",
        str(args.epochs),
        "--batch-size",
        str(args.batch_size),
    ]
    LOG.info("invoking: %s", " ".join(cmd))
    rc = subprocess.call(cmd)
    if rc != 0:
        LOG.error("student fine-tune failed with rc=%d", rc)
        return rc
    LOG.info("distillation complete. promote with scripts/promote_to_ollama.sh %s", args.out / "student")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))