| """Greedy generation with Qwen3.5 (+ optional LoRA adapter) for recognition or translation. |
| |
| Input jsonl rows: |
| ocr: {"id", "image", ...} -> writes {"id", "pred", ...} |
| mt : {"id", <src_field>, "lang", ...} -> writes {"id", "pred", ...} |
| Resumable: ids already present in --out are skipped. Rows are written after every batch. |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from qwen_common import clean_generation, encode_kwargs, load_image, load_model, prompt_text, stop_ids |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--task", choices=["mt", "ocr"], required=True) |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--adapter", default="") |
| ap.add_argument("--input", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--src_field", default="src") |
| ap.add_argument("--image_root", default="") |
| ap.add_argument("--bs", type=int, default=16) |
| ap.add_argument("--max_new_tokens", type=int, default=160) |
| ap.add_argument("--limit", type=int, default=0) |
| ap.add_argument("--tags", action="store_true", help="mt: prompt with <runic-language> <2target> tags (field rl)") |
| ap.add_argument("--force_rl", default="", help="with --tags: use this runic-language tag for every row") |
| args = ap.parse_args() |
| print("ARGS", vars(args), flush=True) |
|
|
| rows = [json.loads(l) for l in open(args.input) if l.strip()] |
| if args.limit: |
| rows = rows[: args.limit] |
| done = set() |
| if os.path.exists(args.out): |
| done = {json.loads(l)["id"] for l in open(args.out) if l.strip()} |
| todo = [r for r in rows if r["id"] not in done] |
| print(f"rows {len(rows)} done {len(done)} todo {len(todo)}", flush=True) |
| if not todo: |
| return |
|
|
| processor, model = load_model(args.model, adapter=args.adapter or None) |
| model.to("cuda").eval() |
| tok = processor.tokenizer |
| tok.padding_side = "left" |
| eos = stop_ids(tok) |
| os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True) |
| fout = open(args.out, "a") |
| for b in range(0, len(todo), args.bs): |
| chunk = todo[b : b + args.bs] |
| if args.task == "ocr": |
| texts = [prompt_text(processor, "ocr") for _ in chunk] |
| images = [load_image(os.path.join(args.image_root, r["image"])) for r in chunk] |
| enc = processor(text=texts, images=images, return_tensors="pt", padding=True) |
| else: |
| texts = [prompt_text(processor, "mt", r[args.src_field] or "", r["lang"], |
| (args.force_rl or r["rl"]) if args.tags else None) for r in chunk] |
| enc = processor(text=texts, return_tensors="pt", padding=True, **encode_kwargs(processor)) |
| enc = {k: v.to("cuda") for k, v in enc.items()} |
| with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): |
| out = model.generate(**enc, max_new_tokens=args.max_new_tokens, do_sample=False, |
| eos_token_id=eos, pad_token_id=tok.pad_token_id, repetition_penalty=1.0) |
| gen = out[:, enc["input_ids"].shape[1]:] |
| for r, g in zip(chunk, gen): |
| pred = clean_generation(tok.decode(g, skip_special_tokens=False)) |
| rec = dict(r) |
| rec["pred"] = pred |
| fout.write(json.dumps(rec, ensure_ascii=False) + "\n") |
| fout.flush() |
| print(f"{min(b + args.bs, len(todo))}/{len(todo)}", flush=True) |
| print("DONE", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
