inference_tsqa.py

#!/usr/bin/env python
"""Evaluate a native ITFormer checkpoint with the unified Time-MQA evaluator."""

from __future__ import annotations

import argparse
import json
import math
import os
import random
import sys
from contextlib import nullcontext
from pathlib import Path

import numpy as np
import torch
from accelerate import Accelerator
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import AutoTokenizer

MQA_DIR = Path(os.environ.get("MQA_DIR", "/mnt/share01/sqk/MQA"))
sys.path.insert(0, str(MQA_DIR))

from data_utils import compute_group_metrics  # noqa: E402
from dataset.tsqa_dataset import ITFormerTSQACollator, ITFormerTSQADataset
from models.TimeLanguageModel import TLM, TLMConfig
from utils.accelerate_compat import patch_accelerate_unwrap_model


def parse_args():
    parser = argparse.ArgumentParser(description="Evaluate ITFormer on Time-MQA TSQA.")
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--eval_path", required=True)
    parser.add_argument(
        "--llm_model_path",
        default="/mnt/share01/sqk/models/qwen2.5-7b-instruct",
    )
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--max_eval_samples", type=int, default=0)
    parser.add_argument("--seed", type=int, default=42)

    parser.add_argument("--d_model", type=int, default=512)
    parser.add_argument("--n_heads", type=int, default=8)
    parser.add_argument("--e_layers", type=int, default=4)
    parser.add_argument("--patch_len", type=int, default=60)
    parser.add_argument("--stride", type=int, default=60)
    parser.add_argument("--input_len", type=int, default=600)
    parser.add_argument("--dropout", type=float, default=0.1)
    parser.add_argument("--it_d_model", type=int, default=896)
    parser.add_argument("--it_n_heads", type=int, default=16)
    parser.add_argument("--it_layers", type=int, default=2)
    parser.add_argument("--it_dropout", type=float, default=0.1)
    parser.add_argument("--itformer_legacy_double_residual", action="store_true")
    parser.add_argument("--prefix_num", type=int, default=25)
    parser.add_argument("--adapter_type", default="itformer", choices=["itformer"])
    parser.add_argument(
        "--llm_attn_implementation",
        default=None,
        choices=["eager", "sdpa", "flash_attention_2"],
    )
    parser.add_argument(
        "--llm_torch_dtype",
        default=None,
        choices=["float16", "bfloat16", "float32"],
    )
    parser.add_argument("--use_lora", action="store_true")
    parser.add_argument("--lora_r", type=int, default=16)
    parser.add_argument("--lora_alpha", type=int, default=32)
    parser.add_argument("--lora_dropout", type=float, default=0.05)
    parser.add_argument(
        "--lora_target_modules",
        nargs="+",
        default=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
    )

    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--num_workers", type=int, default=2)
    parser.add_argument("--max_new_tokens", type=int, default=256)
    parser.add_argument("--bf16", action="store_true")
    return parser.parse_args()


def response_only(text: str) -> str:
    for marker in ("assistant\n", "<|im_start|>assistant\n"):
        if marker in text:
            text = text.rsplit(marker, 1)[-1]
    return text.strip()


def gather_results(local_results, accelerator):
    if accelerator.num_processes == 1:
        return local_results
    from accelerate.utils import gather_object

    gathered = gather_object(local_results)
    if gathered and isinstance(gathered[0], list):
        return [item for part in gathered for item in part]
    return gathered


def main():
    patch_accelerate_unwrap_model()
    args = parse_args()
    accelerator = Accelerator()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    tokenizer_path = (
        args.checkpoint
        if (Path(args.checkpoint) / "tokenizer.json").is_file()
        else args.llm_model_path
    )
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    if "<|image_pad|>" not in tokenizer.get_vocab():
        tokenizer.add_tokens(["<|image_pad|>"])

    tlm_config = TLMConfig(
        llm_model_path=args.llm_model_path,
        freeze_ts_model=True,
        ts_pad_num=args.prefix_num,
        llm_attn_implementation=args.llm_attn_implementation,
        llm_torch_dtype=args.llm_torch_dtype,
        use_lora=args.use_lora,
        lora_r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        lora_target_modules=args.lora_target_modules,
        gradient_checkpointing=False,
    )
    for name in (
        "d_model",
        "n_heads",
        "e_layers",
        "patch_len",
        "stride",
        "input_len",
        "dropout",
        "it_d_model",
        "it_n_heads",
        "it_layers",
        "it_dropout",
        "itformer_legacy_double_residual",
        "prefix_num",
        "adapter_type",
    ):
        setattr(tlm_config, name, getattr(args, name))

    model = TLM.from_pretrained(
        args.checkpoint,
        config=tlm_config,
        ts_config=tlm_config,
    )
    if len(tokenizer) != model.llm_model.config.vocab_size:
        model.llm_model.resize_token_embeddings(len(tokenizer))
    model = accelerator.prepare(model)
    model.eval()

    dataset = ITFormerTSQADataset(
        args.eval_path,
        tokenizer,
        prefix_num=args.prefix_num,
        input_len=args.input_len,
        max_samples=args.max_eval_samples,
        include_answer=False,
    )
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=ITFormerTSQACollator(tokenizer),
        num_workers=args.num_workers,
    )
    loader = accelerator.prepare(loader)

    progress = tqdm(
        loader,
        total=math.ceil(len(dataset) / (args.batch_size * accelerator.num_processes)),
        desc="ITFormer TSQA eval",
        disable=not accelerator.is_main_process,
    )
    results = []
    for batch in progress:
        amp = (
            torch.autocast(device_type="cuda", dtype=torch.bfloat16)
            if args.bf16 and torch.cuda.is_available()
            else nullcontext()
        )
        with torch.no_grad(), amp:
            generated = accelerator.unwrap_model(model).generate(
                input_ids=batch["input_ids"],
                query_ids=batch["query_ids"],
                ts_values=batch["ts_values"],
                stage=batch["stage"],
                attention_mask=batch["attention_mask"],
                max_new_tokens=args.max_new_tokens,
                do_sample=False,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.pad_token_id,
                use_cache=True,
                num_beams=1,
                return_dict_in_generate=False,
            )
        decoded = tokenizer.batch_decode(generated, skip_special_tokens=True)
        for index, text in enumerate(decoded):
            results.append(
                {
                    "_index": int(batch["index"][index].item()),
                    "id": batch["id"][index],
                    "figure_path": batch["figure_path"][index],
                    "application_domain": batch["application_domain"][index],
                    "task_type": batch["task_type"][index],
                    "source_type": batch["source_type"][index],
                    "question_format": batch["question_format"][index],
                    "question": batch["question"][index],
                    "answer": batch["answer"][index],
                    "prediction": response_only(text),
                }
            )

    results = gather_results(results, accelerator)
    if accelerator.is_main_process:
        unique = {}
        for row in results:
            unique[row["_index"]] = row
        results = [unique[index] for index in sorted(unique)]
        for row in results:
            row.pop("_index", None)

        metrics = {"by_group": compute_group_metrics(results)}
        output_dir = Path(args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        with (output_dir / "predictions.jsonl").open("w", encoding="utf-8") as handle:
            for row in results:
                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
        with (output_dir / "metrics.json").open("w", encoding="utf-8") as handle:
            json.dump(metrics, handle, ensure_ascii=False, indent=2)
        print(json.dumps(metrics, ensure_ascii=False, indent=2))
        print(f"Predictions: {output_dir / 'predictions.jsonl'}")
        print(f"Metrics: {output_dir / 'metrics.json'}")

    accelerator.wait_for_everyone()


if __name__ == "__main__":
    main()