Spaces:
Sleeping
Sleeping
File size: 3,339 Bytes
32387a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Run a model on an eval JSONL and save predictions.
Usage on Colab (after training a LoRA adapter)::
python -m scripts.generate \
--eval data/processed/eval_bea_dev.jsonl \
--base-model Qwen/Qwen2.5-3B-Instruct \
--adapter <user>/qwen2.5-3b-gec-sft \
--out results/predictions/sft_bea_dev.jsonl
Predictions JSONL schema:
{
"source": "...", # copied from the eval row
"raw": "...", # exactly what the model produced
"corrected": "...", # parse_inline(raw)[0]
"parse_ok": true/false,
}
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from tqdm import tqdm
from gec.inference import generate_batch, load_model
def read_jsonl(path: Path) -> list[dict]:
with path.open("r", encoding="utf-8") as f:
return [json.loads(line) for line in f if line.strip()]
def write_jsonl(path: Path, rows: list[dict]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", encoding="utf-8") as f:
for row in rows:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--eval", required=True, type=Path,
help="Eval JSONL (must contain a 'source' field).")
ap.add_argument("--base-model", required=True)
ap.add_argument("--adapter", default=None,
help="Optional HF Hub id or local path of a LoRA adapter.")
ap.add_argument("--out", required=True, type=Path)
ap.add_argument("--dtype", default="auto",
choices=["auto", "bfloat16", "float16", "4bit"])
ap.add_argument("--few-shot", action="store_true",
help="Prepend 3-shot examples (use for fair base-model eval).")
ap.add_argument("--max-new-tokens", type=int, default=192)
ap.add_argument("--batch-size", type=int, default=8)
ap.add_argument("--limit", type=int, default=0,
help="Generate only the first N rows (smoke-test).")
args = ap.parse_args()
rows = read_jsonl(args.eval)
if args.limit:
rows = rows[: args.limit]
sentences = [r["source"] for r in rows]
print(f"Loaded {len(sentences)} sources from {args.eval}")
tok, model = load_model(args.base_model, adapter_id=args.adapter, dtype=args.dtype)
print(f"Loaded model {args.base_model}"
+ (f" + adapter {args.adapter}" if args.adapter else "")
+ f" (few_shot={args.few_shot})")
results = []
pbar = tqdm(total=len(sentences), desc="generate")
for start in range(0, len(sentences), args.batch_size):
chunk = sentences[start : start + args.batch_size]
out = generate_batch(
chunk, tok, model,
include_few_shot=args.few_shot,
max_new_tokens=args.max_new_tokens,
batch_size=args.batch_size,
)
for r in out:
results.append({
"source": r.source,
"raw": r.raw,
"corrected": r.corrected,
"parse_ok": r.parse_ok,
})
pbar.update(len(chunk))
pbar.close()
write_jsonl(args.out, results)
print(f"Wrote {len(results)} predictions to {args.out}")
if __name__ == "__main__":
main()
|