quick_distill.py

"""Small-scale DictaBERT distillation — 2K lines, fast completion.

Previous attempts (50K lines) failed due to GPU preemption.
This version uses only 2000 lines → completes in ~20 min.
Even 2K new high-quality examples would help ByT5 improve.
"""

from __future__ import annotations

import json
from pathlib import Path

import modal

datasets_volume = modal.Volume.from_name("rababa-datasets", create_if_missing=True)

image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("build-essential", "git", "curl")
    .pip_install(
        "torch>=2.4,<3",
        "transformers==4.38.0",
        "huggingface_hub>=0.20,<0.25",
        "sentencepiece>=0.2",
        "numpy>=1.26,<3",
        "tqdm>=4.66",
        "pyyaml>=6.0",
    )
    .add_local_dir("src", "/opt/rababa/src", copy=True)
    .workdir("/opt/rababa")
    .env({"PYTHONPATH": "/opt/rababa/src"})
)

app = modal.App(name="rababa", image=image)


@app.function(
    gpu="A10G",
    timeout=30 * 60,
    volumes={"/datasets": datasets_volume},
    secrets=[modal.Secret.from_name("huggingface")],
)
def quick_distill() -> dict:
    """Distill 2000 Hebrew Wikipedia lines with DictaBERT."""
    import torch
    from transformers import AutoModel, AutoTokenizer
    from pathlib import Path as _P

    model_name = "dicta-il/dictabert-large-char-menaked"
    print(f"Loading {model_name}...", flush=True)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
    model.eval()

    # Load Hebrew Wikipedia — only first 2000 lines
    hewiki_path = _P("/datasets/hewiki/train.txt")
    lines = []
    for line in hewiki_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if 10 <= len(line) <= 200:
            lines.append(line)
        if len(lines) >= 2000:
            break

    print(f"Processing {len(lines)} Hebrew Wikipedia lines", flush=True)

    # Batch prediction — 8 at a time
    out_dir = _P("/datasets/hebrew-dictabert-quick")
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / "train.txt"

    batch_size = 8
    count = 0
    with out_path.open("w", encoding="utf-8") as f:
        for i in range(0, len(lines), batch_size):
            batch = lines[i:i + batch_size]
            try:
                predictions = model.predict(batch, tokenizer)
            except Exception as e:
                print(f"Batch {i} error: {e}", flush=True)
                predictions = batch

            for pred in predictions:
                if pred and pred.strip():
                    f.write(pred.strip() + "\n")
                    count += 1

            if i % 200 == 0:
                print(f"  [{i}/{len(lines)}] {count} distilled", flush=True)

    print(f"Distilled {count} lines → {out_path}", flush=True)
    datasets_volume.commit()

    return {"count": count, "path": str(out_path)}


@app.local_entrypoint()
def main():
    result = quick_distill.remote()
    print(json.dumps(result, indent=2, default=str))