Blame for source/scripts/benchmarks/evaluate_external_lm.py

"""Evaluate a pinned Japanese external LM on the four frozen nonblind pools.

Raw joint-tokenized suffix likelihood, without instructions or answer prompts.
The vocabulary head is applied only at scored positions to fit an 8 GiB GPU.
"""
import argparse
import json
import math
from pathlib import Path
import sys
import time

ROOT = Path(__file__).resolve().parents[2]
# Optional packages are isolated from the existing training environment.
OPTIONAL = ROOT / 'artifacts/external-eval/site-packages'
if OPTIONAL.is_dir():
    sys.path.insert(0, str(OPTIONAL))
sys.path.insert(0, str(ROOT / 'src'))

import torch
from torch.nn import functional as F
from vimeml.training.data import file_sha, write_json
from vimeml.training.evaluate_ime import common_prefix_length
from vimeml.benchmarks.evaluate_ajimee import eligibility, metrics, summarize
from vimeml.benchmarks.comparison import paired, rows as read_rows
from vimeml.benchmarks.standard_ime import source_metrics
from vimeml.posttraining.infer import TRACKS, read_snapshot

MODELS = {
    'tinyswallow': {
        'id': 'SakanaAI/TinySwallow-1.5B',
        'revision': '535d1322615c9f8171c3d60b10ef6955cc896fe9',
        'directory': 'artifacts/models/external/TinySwallow-1.5B',
        'output': 'outputs/evaluation-tinyswallow-20261016',
    },
    'rakuten-instruct': {
        'id': 'Rakuten/RakutenAI-2.0-mini-instruct',
        'revision': '6d902489587d324b7d5e201299e4e1a169f3a40b',
        'directory': 'artifacts/models/external/RakutenAI-2.0-mini-instruct',
        'output': 'outputs/evaluation-rakuten-instruct-20261016',
    },
    'sarashina': {
        'id': 'sbintuitions/sarashina2.2-1b',
        'revision': '3bb836ad7475ba192926be66651e4730825df7da',
        'directory': 'artifacts/models/external/sarashina2.2-1b',
        'output': 'outputs/evaluation-sarashina-20261016',
    },
}


class TokenizerAdapter:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def encode(self, text, out_type=int):
        if out_type is not int:
            raise ValueError('Integer token IDs required')
        return self.tokenizer.encode(text, add_special_tokens=False)

    def decode(self, ids):
        return self.tokenizer.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)

    def id_to_piece(self, ident):
        return self.tokenizer.convert_ids_to_tokens(ident)


class ExternalLM:
    def __init__(self, model, tokenizer, candidate_batch=8, head_chunk=128):
        self.model = model.eval()
        self.processor = TokenizerAdapter(tokenizer)
        self.device = next(model.parameters()).device
        self.bos = model.config.bos_token_id
        if type(self.bos) is not int or candidate_batch < 1 or head_chunk < 1:
            raise ValueError('Native BOS and positive evaluation batches required')
        self.candidate_batch, self.head_chunk = candidate_batch, head_chunk
        self.pad = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else self.bos
        self.model.config.context_length = 128

    @torch.inference_mode()
    def score(self, context, candidates):
        if not candidates or any(not c for c in candidates) or len(set(candidates)) != len(candidates):
            raise ValueError('Distinct nonempty candidates required')
        context_ids = [self.bos, *self.processor.encode(context)]
        sequences = [[self.bos, *self.processor.encode(context + c)] for c in candidates]
        common = common_prefix_length([context_ids, *sequences])
        if common < 1 or any(not common < len(s) <= 129 for s in sequences):
            raise ValueError('Invalid suffix or context overflow; no truncation')
        result = []
        for start in range(0, len(sequences), self.candidate_batch):
            group = sequences[start:start + self.candidate_batch]
            width = max(len(s) - 1 for s in group)
            ids = torch.full((len(group), width), self.pad, dtype=torch.long, device=self.device)
            attention = torch.zeros_like(ids)
            row_ids, positions, targets, counts = [], [], [], []
            for row, seq in enumerate(group):
                ids[row, :len(seq) - 1] = torch.tensor(seq[:-1], device=self.device)
                attention[row, :len(seq) - 1] = 1
                suffix = seq[common:]
                counts.append(len(suffix))
                row_ids.extend([row] * len(suffix))
                positions.extend(range(common - 1, len(seq) - 1))
                targets.extend(suffix)
            hidden = self.model.model(input_ids=ids, attention_mask=attention,
                                      use_cache=False, return_dict=True).last_hidden_state
            selected = hidden[torch.tensor(row_ids, device=self.device),
                              torch.tensor(positions, device=self.device)]
            targets = torch.tensor(targets, device=self.device)
            values = []
            for offset in range(0, len(selected), self.head_chunk):
                logits = self.model.lm_head(selected[offset:offset + self.head_chunk])
                logp = -F.cross_entropy(logits.float(), targets[offset:offset + self.head_chunk], reduction='none')
                if not bool(torch.isfinite(logp).all()):
                    raise FloatingPointError('Non-finite full-vocabulary candidate probability')
                values.extend(logp.cpu().tolist())
            offset = 0
            for local, (seq, count) in enumerate(zip(group, counts, strict=True)):
                token_logp = values[offset:offset + count]
                offset += count
                suffix = seq[common:]
                result.append({'text': candidates[start + local],
                    'log_probability_sum': sum(token_logp),
                    'log_probability_mean': sum(token_logp) / count,
                    'scored_tokens': count, 'token_ids': suffix,
                    'pieces': [self.processor.id_to_piece(t) for t in suffix],
                    'token_log_probabilities': token_logp,
                    'boundary_retokenized': seq[:len(context_ids)] != context_ids})
        return {'common_prefix_tokens_including_bos': common,
                'context_tokens_including_bos': len(context_ids), 'candidates': result}


def rerank_external(lm, rows, output, track):
    for position, row in enumerate(rows):
        candidates = row['orders']['azookey']
        row['lm_scores'] = {}
        for field, context, methods in (
            ('contextual', row['left_context'], ('lm_context_sum', 'lm_context_mean_secondary')),
            ('context_free', '', ('lm_no_context_sum',)),
        ):
            reason = eligibility(lm, context, candidates)
            if reason:
                for method in methods:
                    row['orders'][method] = list(candidates)
                    row['fallbacks'][method] = reason
                continue
            # An empty original context gives exactly the same second score.
            scored = row['lm_scores'].get('contextual') if field == 'context_free' and not row['left_context'] else None
            if scored is None:
                scored = lm.score(context, candidates)
            row['lm_scores'][field] = scored
            for method in methods:
                key = 'log_probability_mean' if method.endswith('secondary') else 'log_probability_sum'
                row['orders'][method] = [x['text'] for x in sorted(scored['candidates'], key=lambda x: x[key], reverse=True)]
        if (position + 1) % 25 == 0 or position + 1 == len(rows):
            write_json(output / 'progress.json', {'status': 'running', 'track': track,
                'completed_cases': position + 1, 'total_cases': len(rows)})
            print(f'{track}: {position + 1}/{len(rows)}', flush=True)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model', choices=tuple(MODELS), default='tinyswallow')
    parser.add_argument('--output', type=Path)
    args = parser.parse_args()
    selected = MODELS[args.model]
    model_id, revision = selected['id'], selected['revision']
    model_dir = ROOT / selected['directory']
    output = (args.output or ROOT / selected['output']).resolve()
    if (output / 'plan.json').exists():
        parser.error('An evaluation plan already exists; preserve this run and choose a new output')
    output.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(4)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
        raise ValueError('This frozen external evaluation requires a BF16 CUDA GPU')
    source = json.loads((ROOT / selected['output'] / 'model-source.json').read_text('utf-8'))
    if source['model_id'] != model_id or source['revision'] != revision:
        raise ValueError('Pinned official model source changed')
    files = {name: file_sha(model_dir / name) for name in source['files'] if name != '.gitattributes'}
    frozen = {track: read_snapshot(track) for track in TRACKS}
    inputs = {track: {'cases': len(rows), 'sha256': file_sha(ROOT / f'outputs/evaluation-v4-20261015/inputs/{track}.json')}
              for track, rows in frozen.items()}
    baseline_paths = {'v21': ROOT / 'outputs/evaluation-v4-20261015/v21',
                      'v4b': ROOT / 'outputs/evaluation-v4-b12cosine-20261015/best'}
    if args.model != 'tinyswallow':
        baseline_paths['tinyswallow'] = ROOT / MODELS['tinyswallow']['output']
    baseline_hashes = {name: {track: file_sha(folder / track / 'scores.jsonl') for track in TRACKS}
                       for name, folder in baseline_paths.items()}
    native = json.loads((model_dir / 'config.json').read_text('utf-8'))
    plan = {'model_id': model_id, 'revision': revision, 'model_files_sha256': files,
        'tracks': inputs, 'baseline_score_sha256': baseline_hashes,
        'precision': 'BF16 weights/forward; FP32 full-vocabulary CE; no quantization',
        'scoring': 'raw context+candidate joint tokenization; common-prefix suffix logP sum; native BOS; no final EOS',
        'primary': 'lm_context_sum', 'secondary': ['lm_context_mean_secondary', 'lm_no_context_sum'],
        'context_limit': 128, 'native_context_limit': native['max_position_embeddings'], 'candidate_batch': 8, 'head_chunk': 128,
        'instructions': False, 'chat_template': False, 'reading_conditioning': False,
        'stable_ties': True, 'overflow': 'whole-case AzooKey fallback; no truncation; full denominator',
        'training': False, 'blind_read': False, 'test_split_read': False,
        'training_overlap': 'External pretraining/distillation overlap with public benchmark text is unknown',
        'statistics': 'Paired exact McNemar descriptive results; four correlated exploratory tracks, no multiplicity-adjusted claim',
        'script_sha256': file_sha(Path(__file__)), 'torch': str(torch.__version__)}
    write_json(output / 'plan.json', plan)
    (output / 'evaluation-script.py').write_bytes(Path(__file__).read_bytes())
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import transformers
    began = time.perf_counter()
    tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True, trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(model_dir, local_files_only=True,
        trust_remote_code=False, use_safetensors=True, dtype=torch.bfloat16,
        attn_implementation='sdpa').to('cuda').eval()
    lm = ExternalLM(model, tokenizer)
    torch.cuda.synchronize()
    metadata = {'model_id': model_id, 'revision': revision, 'parameters': sum(p.numel() for p in model.parameters()),
        'vocab_size': model.config.vocab_size, 'native_bos': lm.bos,
        'gpu': torch.cuda.get_device_name(), 'torch': str(torch.__version__), 'transformers': transformers.__version__,
        'precision': plan['precision'], 'model_loading_seconds': time.perf_counter() - began,
        'weights_allocated_MiB': torch.cuda.memory_allocated() / 2**20}
    write_json(output / 'environment.json', metadata)
    reports = {}
    for track, rows in frozen.items():
        started = time.perf_counter()
        rerank_external(lm, rows, output, track)
        destination = output / track
        destination.mkdir(parents=True, exist_ok=True)
        scores = destination / 'scores.jsonl'
        scores.write_text(''.join(json.dumps(r, ensure_ascii=False) + '\n' for r in rows), encoding='utf-8')
        report = {'status': 'complete', 'model': metadata, 'input': inputs[track],
                  'metrics': summarize(rows), 'elapsed_seconds': time.perf_counter() - started,
                  'scores_sha256': file_sha(scores)}
        if 'source' in rows[0]:
            report['source_metrics'] = source_metrics(rows, metrics)
        for name, folder in baseline_paths.items():
            comparison = paired(read_rows(folder / track / 'scores.jsonl'), rows)
            comparison['paired'].pop('significant_at_0_05', None)
            write_json(destination / f'paired-{name}.json', comparison)
        write_json(destination / 'metrics.json', report)
        reports[track] = report['metrics']['all']
        print(track, 'TOP1', reports[track]['lm_context_sum']['top1_correct'], '/', len(rows), flush=True)
    result = {'status': 'complete', 'model': metadata, 'tracks': reports,
        'peak_allocated_MiB': torch.cuda.max_memory_allocated() / 2**20,
        'peak_reserved_MiB': torch.cuda.max_memory_reserved() / 2**20,
        'elapsed_including_load_seconds': time.perf_counter() - began,
        'blind_read': False, 'test_split_read': False}
    write_json(output / 'summary.json', result)
    write_json(output / 'progress.json', result)
    print(json.dumps(result, ensure_ascii=False), flush=True)


if __name__ == '__main__':
    main()