ChipYTY/titans_NPC

1"""2Qwen3-4B baseline evaluation on BABILong QA1 (32k).3 4This script evaluates the pretrained Qwen model WITHOUT any training,5using the same chunk-based streaming approach as the Titans training script.6 7Purpose: Establish a baseline to compare with Titans memory-augmented models.8"""9 10import os11import json12import math13import argparse14import logging15from dataclasses import dataclass16from typing import Optional, Dict, List17 18import torch19import torch.nn as nn20import torch.nn.functional as F21from torch.utils.data import Dataset, DataLoader22from tqdm import tqdm23 24logging.basicConfig(25    level=logging.INFO,26    format="%(asctime)s - %(levelname)s - %(message)s"27)28logger = logging.getLogger(__name__)29 30 31@dataclass32class EvalConfig:33    # paths - same as training config34    model_path: str = "/data/huangyifei/huggingface_cache/hub/models--Qwen--Qwen3-4B-Instruct-2507/snapshots/cdbee75f17c01a7cc42f958dc650907174af0554"35    data_path: str = "/data/yty/BABILong/babilong-train-5k-samples/data/qa1/32k.json"36    output_dir: str = "./outputs/qwen_baseline_eval"37 38    # streaming settings - same as training39    chunk_size: int = 819240    max_length: int = 3276841    answer_reserve_tokens: int = 6442 43    # evaluation44    batch_size: int = 1  # use 1 for simplicity in baseline eval45    max_samples: Optional[int] = 500  # same as training default46    print_examples: int = 2047 48    # precision49    bf16: bool = True50    fp16: bool = False51    use_tf32: bool = True52 53    seed: int = 4254 55 56class BABILongDataset(Dataset):57    """Same dataset class as training script for consistency."""58 59    def __init__(60        self,61        data_path: str,62        tokenizer,63        max_length: int = 32768,64        answer_reserve_tokens: int = 64,65        max_samples: Optional[int] = None,66    ):67        self.tokenizer = tokenizer68        self.max_length = max_length69        self.answer_reserve_tokens = answer_reserve_tokens70 71        logger.info(f"Loading dataset: {data_path}")72        with open(data_path, "r") as f:73            self.data = json.load(f)74 75        if max_samples:76            self.data = self.data[:max_samples]77 78        logger.info(f"Dataset size: {len(self.data)}")79 80    def __len__(self):81        return len(self.data)82 83    def __getitem__(self, idx):84        item = self.data[idx]85        text = f"{item['input']}\n\nQuestion: {item['question']}\nAnswer:"86        target = item["target"]87 88        pad_id = self.tokenizer.pad_token_id or 089        reserve = int(self.answer_reserve_tokens)90 91        prompt_ids = self.tokenizer(92            text,93            max_length=max(self.max_length - reserve, 1),94            truncation=True,95            add_special_tokens=True,96            return_tensors="pt",97        ).input_ids.squeeze(0)98 99        answer_ids = self.tokenizer(100            f" {target}",101            add_special_tokens=False,102            return_tensors="pt",103        ).input_ids.squeeze(0)104 105        available = max(self.max_length - prompt_ids.numel(), 0)106        answer_ids = answer_ids[:available]107 108        input_ids = torch.cat([prompt_ids, answer_ids], dim=0)[: self.max_length]109 110        labels = torch.full_like(input_ids, fill_value=-100)111        if answer_ids.numel() > 0:112            start = prompt_ids.numel()113            end = min(start + answer_ids.numel(), labels.numel())114            labels[start:end] = input_ids[start:end]115 116        seq_len = input_ids.numel()117        if seq_len < self.max_length:118            pad_len = self.max_length - seq_len119            input_ids = F.pad(input_ids, (0, pad_len), value=int(pad_id))120            labels = F.pad(labels, (0, pad_len), value=-100)121            attention_mask = torch.cat(122                [torch.ones(seq_len, dtype=torch.long), torch.zeros(pad_len, dtype=torch.long)],123                dim=0,124            )125        else:126            attention_mask = torch.ones(self.max_length, dtype=torch.long)127 128        return {129            "input_ids": input_ids.to(dtype=torch.long),130            "labels": labels.to(dtype=torch.long),131            "attention_mask": attention_mask,132            "target_text": target,  # keep original target for comparison133        }134 135 136def collate_fn(batch):137    # separate target_text from tensor fields138    target_texts = [b.pop("target_text") for b in batch]139    tensor_batch = {k: torch.stack([b[k] for b in batch], dim=0) for k in batch[0].keys()}140    tensor_batch["target_texts"] = target_texts141    return tensor_batch142 143 144class QwenChunkwiseEvaluator:145    """146    Evaluates Qwen model using chunk-wise streaming (same as training).147    148    Key difference from training: NO memory module, just pure Qwen forward pass.149    Each chunk is processed independently with KV cache reset between samples.150    """151 152    def __init__(self, model, tokenizer, config: EvalConfig, device: torch.device):153        self.model = model154        self.tokenizer = tokenizer155        self.config = config156        self.device = device157        self.hidden_size = model.config.hidden_size158 159    def _split_into_chunks(self, seq_len: int, chunk_size: int):160        """Split sequence into chunks, same as training."""161        chunks = []162        for start in range(0, seq_len, chunk_size):163            end = min(start + chunk_size, seq_len)164            chunks.append((start, end))165        return chunks166 167    @torch.no_grad()168    def evaluate_sample(169        self,170        input_ids: torch.Tensor,171        attention_mask: torch.Tensor,172        labels: torch.Tensor,173    ) -> Dict:174        """175        Evaluate a single sample using chunk-wise streaming.176        177        Process:178        1. Split input into chunks179        2. Process each chunk through Qwen (with overlap for next-token prediction)180        3. Collect predictions only for answer tokens (labels != -100)181        4. Compute loss, token accuracy, and EM accuracy182        """183        batch_size, seq_len = input_ids.shape184        chunk_size = self.config.chunk_size185        chunks = self._split_into_chunks(seq_len, chunk_size)186 187        loss_fct_sum = nn.CrossEntropyLoss(reduction="sum")188        total_loss_sum = 0.0189        total_loss_tokens = 0190 191        pred_tokens: List[int] = []192        target_tokens: List[int] = []193 194        for start, end in chunks:195            # Include 1 overlap token for next-token prediction at chunk boundaries196            proc_start = max(0, start - 1)197            chunk_ids = input_ids[:, proc_start:end]198            chunk_labels = labels[:, proc_start:end]199            chunk_mask = attention_mask[:, proc_start:end]200 201            # Forward pass through Qwen202            outputs = self.model(203                input_ids=chunk_ids,204                attention_mask=chunk_mask,205                use_cache=False,206                output_hidden_states=False,207                return_dict=True,208            )209            logits = outputs.logits  # [batch, seq, vocab]210 211            # Compute loss and predictions for answer tokens212            if chunk_labels is not None and (chunk_labels != -100).any():213                # Shift for next-token prediction214                shift_logits = logits[:, :-1, :].contiguous()215                shift_labels = chunk_labels[:, 1:].contiguous()216 217                valid = shift_labels != -100218                if valid.any():219                    valid_logits = shift_logits[valid]220                    valid_targets = shift_labels[valid]221 222                    # Compute loss223                    chunk_loss = loss_fct_sum(valid_logits.float(), valid_targets)224                    total_loss_sum += chunk_loss.item()225                    total_loss_tokens += valid_targets.numel()226 227                    # Collect predictions228                    pred_ids = torch.argmax(valid_logits, dim=-1)229                    pred_tokens.extend(pred_ids.cpu().tolist())230                    target_tokens.extend(valid_targets.cpu().tolist())231 232        # Compute metrics233        if total_loss_tokens > 0:234            avg_loss = total_loss_sum / total_loss_tokens235        else:236            avg_loss = 0.0237 238        # Token accuracy239        if len(pred_tokens) > 0:240            tok_correct = sum(p == t for p, t in zip(pred_tokens, target_tokens))241            tok_acc = tok_correct / len(pred_tokens)242        else:243            tok_acc = 0.0244 245        # EM accuracy (exact match of decoded strings)246        if len(pred_tokens) > 0:247            pred_text = self.tokenizer.decode(pred_tokens, skip_special_tokens=True).strip()248            target_text = self.tokenizer.decode(target_tokens, skip_special_tokens=True).strip()249            em_match = (pred_text == target_text)250        else:251            pred_text = ""252            target_text = ""253            em_match = False254 255        return {256            "loss": avg_loss,257            "tok_acc": tok_acc,258            "em_match": em_match,259            "pred_text": pred_text,260            "target_text": target_text,261            "num_tokens": len(pred_tokens),262        }263 264    @torch.no_grad()265    def evaluate_dataset(self, dataloader: DataLoader, print_examples: int = 10) -> Dict:266        """Evaluate entire dataset."""267        self.model.eval()268 269        total_loss = 0.0270        total_batches = 0271        total_tok_correct = 0272        total_tok_total = 0273        total_em_correct = 0274        total_em_total = 0275        printed = 0276 277        pbar = tqdm(dataloader, desc="Evaluating", dynamic_ncols=True)278        for batch in pbar:279            input_ids = batch["input_ids"].to(self.device)280            attention_mask = batch["attention_mask"].to(self.device)281            labels = batch["labels"].to(self.device)282            target_texts = batch["target_texts"]283 284            # Process each sample in batch285            for i in range(input_ids.shape[0]):286                result = self.evaluate_sample(287                    input_ids[i:i+1],288                    attention_mask[i:i+1],289                    labels[i:i+1],290                )291 292                if result["num_tokens"] > 0:293                    total_loss += result["loss"]294                    total_batches += 1295                    total_tok_correct += int(result["tok_acc"] * result["num_tokens"])296                    total_tok_total += result["num_tokens"]297                    total_em_correct += int(result["em_match"])298                    total_em_total += 1299 300                    # Print examples301                    if printed < print_examples:302                        logger.info(303                            f"[EVAL SAMPLE {printed + 1}] "304                            f"pred={repr(result['pred_text'])} | "305                            f"label={repr(result['target_text'])} | "306                            f"match={result['em_match']}"307                        )308                        printed += 1309 310                # Update progress bar311                if total_em_total > 0:312                    pbar.set_postfix({313                        "em_acc": f"{total_em_correct / total_em_total * 100:.1f}%",314                        "tok_acc": f"{total_tok_correct / max(total_tok_total, 1) * 100:.1f}%",315                    })316 317        # Compute final metrics318        avg_loss = total_loss / max(total_batches, 1)319        tok_acc = total_tok_correct / max(total_tok_total, 1)320        em_acc = total_em_correct / max(total_em_total, 1)321 322        return {323            "loss": avg_loss,324            "tok_acc": tok_acc,325            "em_acc": em_acc,326            "total_samples": total_em_total,327            "total_tokens": total_tok_total,328        }329 330 331def main():332    from transformers import AutoModelForCausalLM, AutoTokenizer333 334    parser = argparse.ArgumentParser(description="Evaluate Qwen baseline on BABILong")335    parser.add_argument("--model_path", type=str, default=None, help="Path to Qwen model")336    parser.add_argument("--data_path", type=str, default=None, help="Path to BABILong data")337    parser.add_argument("--output_dir", type=str, default=None, help="Output directory")338    parser.add_argument("--max_samples", type=int, default=None, help="Max samples to evaluate")339    parser.add_argument("--chunk_size", type=int, default=None, help="Chunk size for streaming")340    parser.add_argument("--batch_size", type=int, default=1, help="Batch size")341    parser.add_argument("--print_examples", type=int, default=20, help="Number of examples to print")342    parser.add_argument("--eval_split", type=str, default="eval", choices=["train", "eval", "all"],343                        help="Which split to evaluate: train (90%), eval (10%), or all")344    args = parser.parse_args()345 346    config = EvalConfig()347    if args.model_path:348        config.model_path = args.model_path349    if args.data_path:350        config.data_path = args.data_path351    if args.output_dir:352        config.output_dir = args.output_dir353    if args.max_samples is not None:354        config.max_samples = args.max_samples355    if args.chunk_size is not None:356        config.chunk_size = args.chunk_size357    if args.batch_size:358        config.batch_size = args.batch_size359    if args.print_examples is not None:360        config.print_examples = args.print_examples361 362    torch.manual_seed(config.seed)363 364    # Device setup365    if torch.cuda.is_available():366        device = torch.device("cuda")367    else:368        device = torch.device("cpu")369 370    # TF32 settings371    if torch.cuda.is_available() and config.use_tf32:372        torch.backends.cuda.matmul.allow_tf32 = True373        torch.backends.cudnn.allow_tf32 = True374        try:375            torch.set_float32_matmul_precision("high")376        except Exception:377            pass378 379    logger.info("=" * 60)380    logger.info("Qwen3-4B Baseline Evaluation (NO TRAINING)")381    logger.info("=" * 60)382    logger.info(f"model_path: {config.model_path}")383    logger.info(f"data_path: {config.data_path}")384    logger.info(f"output_dir: {config.output_dir}")385    logger.info(f"max_samples: {config.max_samples}")386    logger.info(f"max_length: {config.max_length}")387    logger.info(f"chunk_size: {config.chunk_size}")388    logger.info(f"eval_split: {args.eval_split}")389    logger.info("=" * 60)390 391    # Load tokenizer392    logger.info("Loading tokenizer...")393    tokenizer = AutoTokenizer.from_pretrained(config.model_path, trust_remote_code=True)394    if tokenizer.pad_token is None:395        tokenizer.pad_token = tokenizer.eos_token396 397    # Disable flash-attn checks398    try:399        import transformers400        from transformers.utils import import_utils as _import_utils401 402        def _disabled(*args, **kwargs):403            return False404 405        _import_utils.is_flash_attn_2_available = _disabled406        if hasattr(transformers, "utils") and hasattr(transformers.utils, "is_flash_attn_2_available"):407            transformers.utils.is_flash_attn_2_available = _disabled408        _import_utils.is_torchao_available = _disabled409        if hasattr(transformers, "utils") and hasattr(transformers.utils, "is_torchao_available"):410            transformers.utils.is_torchao_available = _disabled411    except Exception as e:412        logger.warning(f"Disable checks failed (ignored): {e}")413 414    # Load model415    logger.info("Loading model...")416    torch_dtype = torch.bfloat16 if config.bf16 else (torch.float16 if config.fp16 else torch.float32)417    model = AutoModelForCausalLM.from_pretrained(418        config.model_path,419        torch_dtype=torch_dtype,420        device_map=None,421        trust_remote_code=True,422        attn_implementation="sdpa",423        low_cpu_mem_usage=True,424    )425    model.to(device)426    model.config.use_cache = False427    model.eval()428    logger.info(f"Model loaded: {model.config.hidden_size} hidden size, {model.config.num_hidden_layers} layers")429 430    # Load dataset431    logger.info("Loading dataset...")432    full_dataset = BABILongDataset(433        config.data_path,434        tokenizer,435        max_length=config.max_length,436        answer_reserve_tokens=config.answer_reserve_tokens,437        max_samples=config.max_samples,438    )439 440    # Split dataset same as training (90% train, 10% eval)441    train_size = int(0.9 * len(full_dataset))442    eval_size = len(full_dataset) - train_size443    train_dataset, eval_dataset = torch.utils.data.random_split(444        full_dataset,445        [train_size, eval_size],446        generator=torch.Generator().manual_seed(config.seed),447    )448 449    # Select which split to evaluate450    if args.eval_split == "train":451        dataset = train_dataset452        split_name = "train"453    elif args.eval_split == "eval":454        dataset = eval_dataset455        split_name = "eval"456    else:  # all457        dataset = full_dataset458        split_name = "all"459 460    logger.info(f"Evaluating on '{split_name}' split: {len(dataset)} samples")461 462    dataloader = DataLoader(463        dataset,464        batch_size=config.batch_size,465        shuffle=False,466        collate_fn=collate_fn,467        num_workers=0,468    )469 470    # Create evaluator471    evaluator = QwenChunkwiseEvaluator(model, tokenizer, config, device)472 473    # Run evaluation474    logger.info("Starting evaluation...")475    results = evaluator.evaluate_dataset(dataloader, print_examples=config.print_examples)476 477    # Print results478    ppl = math.exp(min(20.0, results["loss"]))479    logger.info("=" * 60)480    logger.info("EVALUATION RESULTS (Qwen Baseline - NO TRAINING)")481    logger.info("=" * 60)482    logger.info(f"Split: {split_name}")483    logger.info(f"Total samples: {results['total_samples']}")484    logger.info(f"Total answer tokens: {results['total_tokens']}")485    logger.info(f"Loss: {results['loss']:.4f}")486    logger.info(f"Perplexity: {ppl:.3f}")487    logger.info(f"Token Accuracy: {results['tok_acc'] * 100:.2f}%")488    logger.info(f"EM Accuracy: {results['em_acc'] * 100:.2f}%")489    logger.info("=" * 60)490 491    # Save results492    os.makedirs(config.output_dir, exist_ok=True)493    results_path = os.path.join(config.output_dir, f"baseline_results_{split_name}.json")494    with open(results_path, "w") as f:495        json.dump({496            "split": split_name,497            "total_samples": int(results["total_samples"]),498            "total_tokens": int(results["total_tokens"]),499            "loss": float(results["loss"]),500            "perplexity": float(ppl),501            "tok_acc_pct": float(results["tok_acc"] * 100),502            "em_acc_pct": float(results["em_acc"] * 100),503            "config": {504                "model_path": config.model_path,505                "data_path": config.data_path,506                "max_samples": config.max_samples,507                "max_length": config.max_length,508                "chunk_size": config.chunk_size,509            }510        }, f, indent=2)511    logger.info(f"Results saved to: {results_path}")512 513 514if __name__ == "__main__":515    main()516