import argparse
import json
import logging

import torch
import transformers

import jlens
from revisions import revision

MODEL = "Qwen/Qwen3.6-27B"

parser = argparse.ArgumentParser()
parser.add_argument("worker", type=int)
parser.add_argument("n_workers", type=int)
parser.add_argument("--prompts", default="prompts_1000.json")
parser.add_argument("--limit", type=int, default=None)
parser.add_argument("--prefix", default="w")
parser.add_argument("--dim-batch", type=int, default=32)
args = parser.parse_args()

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")

with open(args.prompts) as f:
    prompts = json.load(f)[args.worker :: args.n_workers]
if args.limit:
    prompts = prompts[: args.limit]
print(f"worker {args.worker}/{args.n_workers}: {len(prompts)} prompts", flush=True)

try:
    hf = transformers.AutoModelForCausalLM.from_pretrained(MODEL, revision=revision(MODEL), dtype=torch.bfloat16)
except (ValueError, KeyError) as exc:
    print(f"AutoModelForCausalLM failed ({exc}); trying AutoModelForImageTextToText", flush=True)
    hf = transformers.AutoModelForImageTextToText.from_pretrained(MODEL, revision=revision(MODEL), dtype=torch.bfloat16)
hf = hf.cuda()
tok = transformers.AutoTokenizer.from_pretrained(MODEL, revision=revision(MODEL))
model = jlens.from_hf(hf, tok)
print(model, flush=True)

lens = jlens.fit(
    model,
    prompts,
    dim_batch=args.dim_batch,
    max_seq_len=128,
    checkpoint_path=f"ckpt_{args.prefix}{args.worker}.pt",
    checkpoint_every=5,
)
lens.save(f"lens_{args.prefix}{args.worker}.pt")
print(
    f"peak GPU mem: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB; "
    f"saved lens_{args.prefix}{args.worker}.pt",
    flush=True,
)