evaluate.py

"""
Evaluation wrapper for ailmo using EleutherAI's lm-evaluation-harness.

Loads a checkpoint and runs standard NLP benchmarks to measure model quality.
This is the industry-standard evaluation framework used by HuggingFace's
Open LLM Leaderboard, NVIDIA, Cohere, and most LLM research papers.

Usage:
    python evaluate.py --checkpoint checkpoints/final.pt
    python evaluate.py --checkpoint checkpoints/final.pt --tasks hellaswag,arc_easy
    python evaluate.py --checkpoint checkpoints/step_5000.pt --batch-size 32

Recommended benchmarks for 100M models:
    - hellaswag:       Commonsense reasoning (random ~25%, good 100M: 26-30%)
    - arc_easy:        Elementary science QA (random ~25%, good 100M: 25-30%)
    - lambada_openai:  Next-word prediction (random ~0%, good 100M: 5-15%)

Expected performance at 100M params: slightly above random baselines.
These benchmarks are informative even at small scale — they show whether
the model has learned language structure beyond memorization.
"""

import argparse
import json
import math
import os
import torch
import tiktoken

from configs import ModelConfig
from model import LLM
from generate import generate as generate_text


class AilmoLM:
    """lm-eval-harness compatible wrapper for the ailmo model.

    Implements the interface expected by lm_eval.simple_evaluate():
    tokenization, log-likelihood computation, and text generation.
    """

    def __init__(self, checkpoint_path: str, device: str = "cuda", batch_size: int = 16):
        self.device = torch.device(device)
        self._batch_size = batch_size

        # Load checkpoint
        print(f"Loading checkpoint: {checkpoint_path}")
        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        self.config = ckpt["model_config"]

        self.model = LLM(self.config).to(self.device)
        self.model.load_state_dict(ckpt["model_state_dict"])
        self.model.eval()
        print(f"Model loaded: {self.model.count_parameters():,} params")

        # Tokenizer
        self.enc = tiktoken.get_encoding("gpt2")
        self.eot_token = self.enc.eot_token  # end of text token

    def encode(self, text: str) -> list[int]:
        return self.enc.encode(text, allowed_special=set())

    def decode(self, tokens: list[int]) -> str:
        return self.enc.decode(tokens)

    @torch.no_grad()
    def loglikelihood(self, context: str, continuation: str) -> tuple[float, bool]:
        """Compute log-likelihood of continuation given context.

        This is the core operation for multiple-choice benchmarks like
        HellaSwag and ARC. For each answer option, we compute how likely
        the model thinks that continuation is given the context.

        Returns:
            (total_log_prob, is_greedy): total log probability of the
            continuation tokens, and whether each token was the argmax.
        """
        ctx_tokens = self.encode(context)
        cont_tokens = self.encode(continuation)
        all_tokens = ctx_tokens + cont_tokens

        # Truncate from the left if too long
        if len(all_tokens) > self.config.context_length:
            all_tokens = all_tokens[-self.config.context_length:]
            # Adjust continuation length
            cont_len = min(len(cont_tokens), len(all_tokens))
        else:
            cont_len = len(cont_tokens)

        # Model input: all tokens except the last (we predict the last from second-to-last)
        input_ids = torch.tensor([all_tokens[:-1]], device=self.device)
        target_ids = torch.tensor([all_tokens[1:]], device=self.device)

        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits = self.model(input_ids)

        # Log probabilities
        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)

        # We only care about the continuation portion
        # cont_start is the index in the output where continuation begins
        cont_start = len(all_tokens) - cont_len - 1
        cont_log_probs = log_probs[0, cont_start:, :]
        cont_targets = target_ids[0, cont_start:]

        # Gather log probs of the actual continuation tokens
        token_log_probs = cont_log_probs.gather(1, cont_targets.unsqueeze(1)).squeeze(1)
        total_log_prob = token_log_probs.sum().item()

        # Check if continuation tokens are greedy (argmax at each position)
        is_greedy = (cont_log_probs.argmax(dim=-1) == cont_targets).all().item()

        return total_log_prob, bool(is_greedy)

    @torch.no_grad()
    def loglikelihood_rolling(self, text: str) -> float:
        """Compute rolling log-likelihood for a full text.

        Used by perplexity-based benchmarks like LAMBADA.
        Processes the text in context-length chunks.
        """
        tokens = self.encode(text)
        total_ll = 0.0
        ctx_len = self.config.context_length

        # Process in chunks
        for start in range(0, len(tokens) - 1, ctx_len):
            chunk = tokens[start:start + ctx_len + 1]
            if len(chunk) < 2:
                break

            input_ids = torch.tensor([chunk[:-1]], device=self.device)
            target_ids = torch.tensor([chunk[1:]], device=self.device)

            with torch.autocast("cuda", dtype=torch.bfloat16):
                logits = self.model(input_ids)

            log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
            token_log_probs = log_probs[0].gather(1, target_ids[0].unsqueeze(1)).squeeze(1)
            total_ll += token_log_probs.sum().item()

        return total_ll

    def generate(self, context: str, max_tokens: int = 128, temperature: float = 0.0) -> str:
        """Generate text continuation (greedy by default for eval)."""
        return generate_text(
            self.model, context, max_tokens, temperature,
            top_k=0 if temperature == 0 else 50,
            device=self.device,
        )


def run_lm_eval(model_wrapper: AilmoLM, tasks: list[str], output_dir: str, num_fewshot: int = 0):
    """Run lm-eval-harness benchmarks using our model."""
    try:
        import lm_eval
        from lm_eval.api.model import LM
        from lm_eval.api.instance import Instance
    except ImportError:
        print("Error: lm-eval not installed. Run: pip install lm-eval")
        print("\nFalling back to built-in evaluation...")
        return run_builtin_eval(model_wrapper, output_dir)

    # Create an lm-eval compatible adapter
    class AilmoLMEval(LM):
        def __init__(self, wrapper):
            super().__init__()
            self._wrapper = wrapper

        @property
        def eot_token_id(self):
            return self._wrapper.eot_token

        @property
        def max_length(self):
            return self._wrapper.config.context_length

        @property
        def max_gen_toks(self):
            return 256

        @property
        def batch_size(self):
            return self._wrapper._batch_size

        @property
        def device(self):
            return self._wrapper.device

        def tok_encode(self, string, **kwargs):
            return self._wrapper.encode(string)

        def tok_decode(self, tokens, **kwargs):
            return self._wrapper.decode(tokens)

        def _model_generate(self, context, max_length, stop, **kwargs):
            # Not typically used for multiple-choice tasks
            text = self._wrapper.generate(
                self._wrapper.decode(context[0].tolist()),
                max_tokens=max_length,
            )
            return [self._wrapper.encode(text)]

        def loglikelihood(self, requests):
            results = []
            for req in requests:
                ctx, cont = req.args
                ll, is_greedy = self._wrapper.loglikelihood(ctx, cont)
                results.append((ll, is_greedy))
            return results

        def loglikelihood_rolling(self, requests):
            results = []
            for req in requests:
                text = req.args[0]
                ll = self._wrapper.loglikelihood_rolling(text)
                results.append((ll,))
            return results

        def generate_until(self, requests):
            results = []
            for req in requests:
                ctx = req.args[0]
                gen_kwargs = req.args[1] if len(req.args) > 1 else {}
                max_tokens = gen_kwargs.get("max_gen_toks", 128)
                text = self._wrapper.generate(ctx, max_tokens=max_tokens, temperature=0.0)
                # Strip the context from the output
                if text.startswith(ctx):
                    text = text[len(ctx):]
                results.append(text)
            return results

    adapter = AilmoLMEval(model_wrapper)

    print(f"\nRunning lm-eval benchmarks: {', '.join(tasks)}")
    print(f"  Few-shot: {num_fewshot}")

    results = lm_eval.simple_evaluate(
        model=adapter,
        tasks=tasks,
        num_fewshot=num_fewshot,
        batch_size=model_wrapper._batch_size,
    )

    # Save results
    os.makedirs(output_dir, exist_ok=True)

    # Print results table
    print("\n" + "=" * 60)
    print("  Benchmark Results")
    print("=" * 60)
    for task_name, task_results in results["results"].items():
        print(f"\n  {task_name}:")
        for metric, value in task_results.items():
            if isinstance(value, float):
                print(f"    {metric}: {value:.4f}")
            else:
                print(f"    {metric}: {value}")

    return results["results"]


def run_builtin_eval(model_wrapper: AilmoLM, output_dir: str) -> dict:
    """Fallback evaluation using simple perplexity on validation data."""
    from data import TokenDataset
    import numpy as np

    print("\nRunning built-in perplexity evaluation...")

    val_path = "data/val.bin"
    if not os.path.exists(val_path):
        print(f"  No validation data at {val_path}")
        return {}

    data = np.memmap(val_path, dtype=np.uint16, mode="r")
    ctx_len = model_wrapper.config.context_length

    total_ll = 0.0
    total_tokens = 0
    n_chunks = min(50, len(data) // (ctx_len + 1))

    for i in range(n_chunks):
        start = i * ctx_len
        chunk = torch.tensor(data[start:start + ctx_len + 1].astype(int), device=model_wrapper.device).unsqueeze(0)
        input_ids = chunk[:, :-1]
        target_ids = chunk[:, 1:]

        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            logits = model_wrapper.model(input_ids)
        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
        token_ll = log_probs.gather(2, target_ids.unsqueeze(2)).squeeze(2)
        total_ll += token_ll.sum().item()
        total_tokens += target_ids.numel()

    avg_ll = total_ll / total_tokens
    perplexity = math.exp(-avg_ll)

    print(f"\n  Validation perplexity: {perplexity:.2f}")
    print(f"  Avg log-likelihood: {avg_ll:.4f}")
    print(f"  Tokens evaluated: {total_tokens:,}")

    results = {
        "perplexity": round(perplexity, 2),
        "avg_log_likelihood": round(avg_ll, 4),
        "tokens_evaluated": total_tokens,
    }
    return results


def main():
    parser = argparse.ArgumentParser(description="Evaluate ailmo model")
    parser.add_argument("--checkpoint", type=str, required=True,
                        help="Path to model checkpoint")
    parser.add_argument("--tasks", type=str, default="hellaswag,arc_easy,lambada_openai",
                        help="Comma-separated benchmark tasks")
    parser.add_argument("--output", type=str, default="results/eval_results",
                        help="Output directory for results")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--num-fewshot", type=int, default=0,
                        help="Number of few-shot examples (default: 0)")
    args = parser.parse_args()

    # Load model
    model_wrapper = AilmoLM(args.checkpoint, batch_size=args.batch_size)

    # Run evaluation
    tasks = [t.strip() for t in args.tasks.split(",")]
    results = run_lm_eval(model_wrapper, tasks, args.output, args.num_fewshot)

    # Save results JSON
    os.makedirs(args.output, exist_ok=True)
    ckpt_name = os.path.basename(args.checkpoint).replace(".pt", "")
    output_path = os.path.join(args.output, f"{ckpt_name}_eval.json")
    with open(output_path, "w") as f:
        json.dump({
            "checkpoint": args.checkpoint,
            "tasks": tasks,
            "num_fewshot": args.num_fewshot,
            "results": results,
        }, f, indent=2)
    print(f"\nResults saved: {output_path}")


if __name__ == "__main__":
    main()