| """LoRA fine-tuning of Qwen3.5 for runic recognition (image -> transliteration) or |
| translation (transliteration -> English/Swedish). Loss is computed on the answer tokens only. |
| |
| Data (jsonl): |
| mt : {"id", "src", "tgt", "lang"} |
| ocr: {"id", "image", "target"} |
| """ |
| import argparse |
| import json |
| import math |
| import os |
| import random |
| import sys |
| import time |
|
|
| import torch |
| from torch.utils.data import DataLoader, Dataset |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from ocr_noise import corrupt |
| from qwen_common import encode_kwargs, end_text, load_image, load_model, prompt_text |
|
|
|
|
| def read_jsonl(paths, limit=None): |
| rows = [] |
| for p in paths.split(","): |
| with open(p) as f: |
| rows.extend(json.loads(l) for l in f if l.strip()) |
| if limit: |
| rows = rows[:limit] |
| return rows |
|
|
|
|
| class Examples(Dataset): |
| def __init__(self, rows, processor, task, image_root, noise_p=0.0, noise_max=0.5, seed=0, tags=False): |
| self.rows, self.processor, self.task, self.tags = rows, processor, task, tags |
| self.image_root, self.noise_p, self.noise_max = image_root, noise_p, noise_max |
| self.rng = random.Random(seed) |
| self.end = end_text(processor.tokenizer) |
| self.kw = encode_kwargs(processor) |
|
|
| def __len__(self): |
| return len(self.rows) |
|
|
| def __getitem__(self, i): |
| r = self.rows[i] |
| if self.task == "ocr": |
| img = load_image(os.path.join(self.image_root, r["image"])) |
| prompt = prompt_text(self.processor, "ocr") |
| target = r["target"] |
| full = self.processor(text=[prompt + target + self.end], images=[img], return_tensors="pt") |
| n_prompt = self.processor(text=[prompt], images=[img], return_tensors="pt")["input_ids"].shape[1] |
| else: |
| src = r["src"] |
| rng = random.Random(int.from_bytes(os.urandom(8), "little")) |
| if self.noise_p > 0 and rng.random() < self.noise_p: |
| src = corrupt(src, rng.uniform(0.05, self.noise_max), rng) |
| prompt = prompt_text(self.processor, "mt", src, r["lang"], r["rl"] if self.tags else None) |
| target = r["tgt"] |
| full = self.processor(text=[prompt + target + self.end], return_tensors="pt", **self.kw) |
| n_prompt = self.processor(text=[prompt], return_tensors="pt", **self.kw)["input_ids"].shape[1] |
| item = {k: v for k, v in full.items()} |
| labels = item["input_ids"].clone() |
| labels[:, :n_prompt] = -100 |
| item["labels"] = labels |
| return item |
|
|
|
|
| def collate(batch, pad_id, max_len): |
| out = {} |
| seq_keys = [k for k, v in batch[0].items() if v.dim() == 2 and v.shape[0] == 1 and k != "image_grid_thw"] |
| L = min(max(b["input_ids"].shape[1] for b in batch), max_len) |
| for k in seq_keys: |
| pad = pad_id if k == "input_ids" else (-100 if k == "labels" else 0) |
| rows = [] |
| for b in batch: |
| t = b[k][0, :L] |
| rows.append(torch.cat([t, torch.full((L - t.shape[0],), pad, dtype=t.dtype)])) |
| out[k] = torch.stack(rows) |
| for k in batch[0]: |
| if k not in seq_keys: |
| out[k] = torch.cat([b[k] for b in batch], dim=0) |
| return out |
|
|
|
|
| @torch.no_grad() |
| def val_loss(model, loader, device, max_batches=50): |
| model.eval() |
| tot, n = 0.0, 0 |
| for bi, batch in enumerate(loader): |
| if bi >= max_batches: |
| break |
| batch = {k: v.to(device) for k, v in batch.items()} |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| tot += model(**batch).loss.item() |
| n += 1 |
| model.train() |
| return tot / max(n, 1) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--task", choices=["mt", "ocr"], required=True) |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--train_file", required=True) |
| ap.add_argument("--val_file", required=True) |
| ap.add_argument("--image_root", default="") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--epochs", type=float, default=3) |
| ap.add_argument("--lr", type=float, default=2e-4) |
| ap.add_argument("--bs", type=int, default=8) |
| ap.add_argument("--grad_accum", type=int, default=2) |
| ap.add_argument("--rank", type=int, default=16) |
| ap.add_argument("--alpha", type=int, default=32) |
| ap.add_argument("--max_len", type=int, default=1536) |
| ap.add_argument("--max_steps", type=int, default=0) |
| ap.add_argument("--limit", type=int, default=0) |
| ap.add_argument("--noise_p", type=float, default=0.0) |
| ap.add_argument("--seed", type=int, default=13) |
| ap.add_argument("--workers", type=int, default=6) |
| ap.add_argument("--init_adapter", default="") |
| ap.add_argument("--tags", action="store_true", help="mt: prepend <runic-language> <2target> tags (field rl)") |
| ap.add_argument("--no_last", action="store_true", help="do not save last_adapter (disk)") |
| args = ap.parse_args() |
| print("ARGS", vars(args), flush=True) |
| torch.manual_seed(args.seed) |
| random.seed(args.seed) |
| os.makedirs(args.out, exist_ok=True) |
|
|
| processor, model = load_model(args.model, adapter=args.init_adapter or None, trainable=bool(args.init_adapter)) |
| tok = processor.tokenizer |
| if not args.init_adapter: |
| from peft import LoraConfig, get_peft_model |
|
|
| cfg = LoraConfig(r=args.rank, lora_alpha=args.alpha, lora_dropout=0.05, target_modules="all-linear", |
| task_type="CAUSAL_LM") |
| model = get_peft_model(model, cfg) |
| model.print_trainable_parameters() |
| print("base dtype", next(p for n, p in model.named_parameters() if "lora" not in n).dtype, |
| "lora dtype", next(p for n, p in model.named_parameters() if "lora" in n).dtype, flush=True) |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| model.enable_input_require_grads() |
| device = "cuda" |
| model.to(device) |
| model.train() |
|
|
| tr_rows = read_jsonl(args.train_file, args.limit or None) |
| va_rows = read_jsonl(args.val_file, 400 if not args.limit else min(args.limit, 40)) |
| random.Random(args.seed).shuffle(tr_rows) |
| tr = Examples(tr_rows, processor, args.task, args.image_root, args.noise_p, seed=args.seed, tags=args.tags) |
| va = Examples(va_rows, processor, args.task, args.image_root, 0.0, seed=args.seed, tags=args.tags) |
| col = lambda b: collate(b, tok.pad_token_id, args.max_len) |
| tl = DataLoader(tr, batch_size=args.bs, shuffle=True, num_workers=args.workers, collate_fn=col, |
| persistent_workers=args.workers > 0) |
| vl = DataLoader(va, batch_size=args.bs, shuffle=False, num_workers=args.workers, collate_fn=col) |
|
|
| params = [p for p in model.parameters() if p.requires_grad] |
| opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0) |
| steps_per_epoch = math.ceil(len(tl) / args.grad_accum) |
| total = args.max_steps or int(steps_per_epoch * args.epochs) |
| warm = max(1, int(0.03 * total)) |
| sched = torch.optim.lr_scheduler.LambdaLR( |
| opt, lambda s: min(1.0, (s + 1) / warm) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / max(1, total))))) |
| print(f"train {len(tr)} val {len(va)} steps/epoch {steps_per_epoch} total {total}", flush=True) |
|
|
| step, best, t0 = 0, float("inf"), time.time() |
| log = open(os.path.join(args.out, "train_log.jsonl"), "a") |
| done = False |
| epoch = 0 |
| while not done: |
| for bi, batch in enumerate(tl): |
| batch = {k: v.to(device) for k, v in batch.items()} |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| loss = model(**batch).loss / args.grad_accum |
| loss.backward() |
| if (bi + 1) % args.grad_accum: |
| continue |
| torch.nn.utils.clip_grad_norm_(params, 1.0) |
| opt.step() |
| sched.step() |
| opt.zero_grad(set_to_none=True) |
| step += 1 |
| if step % 20 == 0 or step == 1: |
| rec = dict(step=step, loss=loss.item() * args.grad_accum, lr=sched.get_last_lr()[0], |
| elapsed=round(time.time() - t0)) |
| print(json.dumps(rec), flush=True) |
| log.write(json.dumps(rec) + "\n") |
| log.flush() |
| eval_every = max(1, steps_per_epoch // 2) if not args.max_steps else args.max_steps |
| if step % eval_every == 0 or step >= total: |
| vloss = val_loss(model, vl, device) |
| rec = dict(step=step, epoch=round(step / steps_per_epoch, 2), val_loss=vloss) |
| print(json.dumps(rec), flush=True) |
| log.write(json.dumps(rec) + "\n") |
| log.flush() |
| if vloss < best: |
| best = vloss |
| model.save_pretrained(os.path.join(args.out, "best_adapter")) |
| json.dump(rec, open(os.path.join(args.out, "best.json"), "w")) |
| if step >= total: |
| done = True |
| break |
| epoch += 1 |
| if not args.no_last: |
| model.save_pretrained(os.path.join(args.out, "last_adapter")) |
| json.dump(dict(total_steps=step, best_val_loss=best, minutes=round((time.time() - t0) / 60, 1)), |
| open(os.path.join(args.out, "DONE.json"), "w")) |
| print("DONE", step, best, flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
