| """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 = 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 |
| |
| 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() |
|
|
