#!/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()