"""Reproduce the ARTI English/Chinese Transformer reading-QA comparison. Use the pinned reproduction dependencies in reading-requirements.txt. All downloads, data, weights and runs stay in --work. No browser inference. Data: SQuAD 1.1 and CMRC2018 (CC BY-SA 4.0), original IDs retained. Scoring: official SQuAD answerable EM/F1 core and CMRC2018 Python 3 scorer. """ from __future__ import annotations import argparse import gc import hashlib import importlib.util import json import os from pathlib import Path import random import sys import time import urllib.request import numpy as np import torch from torch import nn from transformers import AutoModelForCausalLM, AutoTokenizer import arti from arti.language_recipe import ( CausalLMRecipe, CausalLMRecipeConfig, CausalLMModelConfig, CausalLMOptimizerConfig, CausalLMGenerationConfig, CausalLMPredictor, ) from arti.recipe_adapters import FormulaAdapterConfig from arti.recipe_chat import encode_final_answer, collate_chat, ChatBatches from arti.training_config import ARTITrainingConfig from arti.training_run import ARTITrainingRunConfig BASE = 'Qwen/Qwen3-0.6B' REVISION = 'c1899de289a04d12100db370d81485cdf75e47ca' SQUAD_REV = 'eee5fdbf62f8613a7812b03419e6b29617b74fd1' CMRC_REV = 'c0eb1b6ba219847457e6af3180da722bbeb656af' TEMPLATE = {'enable_thinking': False} SEEDS = (17, 29, 43) METHODS = ('base', 'torch', 'peft', 'arti') PATHS = ('model.layers.26.mlp', 'model.layers.27.mlp') def digest(value): return hashlib.sha256(value if isinstance(value, bytes) else str(value).encode()).hexdigest() def write(path, value): path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + '\n', encoding='utf8') def seed_all(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def downloads(work): sources = work / 'sources' sources.mkdir(parents=True, exist_ok=True) specs = {} for lang, repo, rev, prefix in ( ('squad', 'rajpurkar/SQuAD-explorer', SQUAD_REV, 'dataset'), ('cmrc', 'ymcui/cmrc2018', CMRC_REV, 'squad-style-data'), ): for split in ('train', 'dev'): name = f'{split}-v1.1.json' if lang == 'squad' else f'cmrc2018_{split}.json' specs[f'{lang}-{split}.json'] = f'https://raw.githubusercontent.com/{repo}/{rev}/{prefix}/{name}' specs['squad_evaluate.py'] = f'https://raw.githubusercontent.com/rajpurkar/SQuAD-explorer/{SQUAD_REV}/evaluate-v2.0.py' specs['cmrc_evaluate_v6.py'] = f'https://raw.githubusercontent.com/ymcui/cmrc2018/{CMRC_REV}/baseline/cmrc2018_evaluate.py' receipts = {} for name, url in specs.items(): target = sources / name if not target.exists(): target.write_bytes(urllib.request.urlopen(url).read()) receipts[name] = {'url': url, 'sha256': digest(target.read_bytes()), 'bytes': target.stat().st_size} write(sources / 'manifest.json', receipts) return receipts def messages(lang, context, question, answer=None): instruction = ( 'Read the passage and answer the question. Return only the shortest exact answer from the passage, without explanation.' if lang == 'squad' else '阅读材料并回答问题。只输出材料中最简短、准确的答案，不要解释。' ) body = f'Passage:\n{context}\n\nQuestion: {question}' if lang == 'squad' else f'材料：\n{context}\n\n问题：{question}' result = [{'role': 'system', 'content': instruction}, {'role': 'user', 'content': body}] if answer is not None: result.append({'role': 'assistant', 'content': answer}) return result def prepare(work, base_path): downloads(work) import nltk nltk.data.path.insert(0, str(work / 'nltk-data')) try: nltk.data.find('tokenizers/punkt_tab/english/') except LookupError: if not nltk.download('punkt_tab', download_dir=work / 'nltk-data', quiet=True): raise RuntimeError('CMRC official scoring requires NLTK punkt_tab data') tokenizer = AutoTokenizer.from_pretrained(base_path, local_files_only=True) for lang in ('squad', 'cmrc'): records, stats = {}, {} for split in ('train', 'dev'): source = json.loads((work / 'sources' / f'{lang}-{split}.json').read_text(encoding='utf8')) eligible = [] seen, too_long, answer_long = 0, 0, 0 for article in source['data']: for paragraph in article['paragraphs']: context = paragraph['context'] context_hash = digest(context) # Necessary prefilter; every retained row still gets its complete native template. context_over = len(tokenizer.encode(context, add_special_tokens=False)) >= 384 for qa in paragraph['qas']: seen += 1 refs = list(dict.fromkeys(a['text'] for a in qa['answers'])) answer = refs[0] if context_over: too_long += 1 continue if len(tokenizer.encode(answer, add_special_tokens=False)) > 32: answer_long += 1 continue turns = messages(lang, context, qa['question'], answer) ex = encode_final_answer(qa['id'], turns, tokenizer, template_kwargs=TEMPLATE) if len(ex.input_ids) > 384: too_long += 1 continue assert ex.labels[:ex.prompt_length].eq(-100).all() assert ex.labels[-1].item() == tokenizer.eos_token_id eligible.append({'id': qa['id'], 'context_sha256': context_hash, 'context': context, 'question': qa['question'], 'references': refs, 'messages': turns, 'tokens': len(ex.input_ids), 'target_tokens': ex.target_tokens}) stats[split] = {'source_rows': seen, 'eligible_rows': len(eligible), 'overlength_excluded': too_long, 'answer_over_32_excluded': answer_long, 'eligible_contexts': len(set(x['context_sha256'] for x in eligible))} # Deterministic, context-first order fixed before any model outcomes. records[split] = sorted(eligible, key=lambda x: (digest('arti-qa-v1:' + x['context_sha256']), digest(x['id']))) validation = records['train'][:128] validation_contexts = {x['context_sha256'] for x in validation} train = [x for x in records['train'] if x['context_sha256'] not in validation_contexts][:512] train_contexts = {x['context_sha256'] for x in train} heldout = [x for x in records['dev'] if x['context_sha256'] not in train_contexts | validation_contexts][:256] assert len(train) == 512 and len(validation) == 128 and len(heldout) == 256, stats assert not train_contexts & validation_contexts all_ids = [x['id'] for x in train + validation + heldout] assert len(all_ids) == len(set(all_ids)) target = work / lang target.mkdir(parents=True, exist_ok=True) for name, rows in [('train', train), ('validation', validation), ('heldout', heldout)]: write(target / f'{name}.json', rows) with (target / 'chat.jsonl').open('w', encoding='utf8') as stream: for split, rows in [('train', train), ('validation', validation)]: for row in rows: stream.write(json.dumps({'id': row['id'], 'split': split, 'messages': row['messages']}, ensure_ascii=False) + '\n') manifest = {'dataset': lang, 'selection': 'sha256 context-first then original-ID order; validation contexts removed from train; public dev heldout', 'counts': {'train': 512, 'validation': 128, 'heldout': 256}, 'eligibility': stats, 'splits': {name: {'ids': [x['id'] for x in rows], 'sha256': digest(json.dumps(rows, ensure_ascii=False))} for name, rows in [('train', train), ('validation', validation), ('heldout', heldout)]}, 'tokenizer': BASE, 'revision': REVISION, 'template_kwargs': TEMPLATE, 'limits': {'complete_chat_tokens': 384, 'answer_tokens': 32}, 'seed_selection': 'none'} write(target / 'data-manifest.json', manifest) print(json.dumps({'prepared': lang, 'eligibility': stats}), flush=True) class TorchResidual(nn.Module): def __init__(self, original, dim=1024, rank=32): super().__init__() self.original = original # Public Qwen attachment first casts the layer to host BF16, then .float(). # Match that initialization round trip; subsequent adapter computation is FP32. self.down = nn.Parameter((torch.randn(dim, rank) / dim ** .5).to(torch.bfloat16).float()) self.up = nn.Parameter(torch.zeros(rank, dim)) def forward(self, x): x = self.original(x) with torch.autocast(device_type=x.device.type, enabled=False): operand = x.float() value = operand + torch.nn.functional.silu(operand @ self.down) @ self.up return value.to(x.dtype) def load_base(base_path): return AutoModelForCausalLM.from_pretrained(base_path, local_files_only=True, dtype=torch.bfloat16, attn_implementation='sdpa').to('cuda') def config(work, lang, seed, output): return CausalLMRecipeConfig( data_path=work / lang / 'chat.jsonl', model=CausalLMModelConfig(id=BASE, revision=REVISION, dtype='bfloat16', local_files_only=True), adapter=FormulaAdapterConfig(rank=32, last_layers=2), batch_size=8, max_length=384, overlength='error', template_kwargs=TEMPLATE, seed=seed, group_by_length=True, max_grad_norm=1., training=ARTITrainingConfig(device='cuda', steps=200, learning_rate=.001, mixed_precision='no'), optimizer=CausalLMOptimizerConfig(weight_decay=.01, implementation='fused'), generation=CausalLMGenerationConfig(max_new_tokens=32, batch_size=8), run=ARTITrainingRunConfig(output_dir=output, checkpoint_steps=200, checkpoint_seconds=86400., log_steps=20), cache_dir=work / lang / 'prepared-cache') def model_for(method, base_path, work, lang, seed, output): seed_all(seed) model = load_base(base_path) tokenizer = AutoTokenizer.from_pretrained(base_path, local_files_only=True) model.requires_grad_(False) model.config.use_cache = False recipe = None if method == 'arti': recipe = CausalLMRecipe(config(work, lang, seed, output / 'native-run'), model=model, tokenizer=tokenizer) model = recipe.model elif method == 'torch': # The same CPU initialization order, function and post-MLP placement as ARTI. torch.manual_seed(seed) for layer in model.model.layers[-2:]: layer.mlp = TorchResidual(layer.mlp).to('cuda') elif method == 'peft': from peft import LoraConfig, get_peft_model torch.manual_seed(seed) model = get_peft_model(model, LoraConfig(r=32, lora_alpha=32, lora_dropout=0., bias='none', target_modules=r'model\.layers\.(26|27)\.self_attn\.v_proj', task_type='CAUSAL_LM')) parameters = [p for p in model.parameters() if p.requires_grad] if method != 'base': assert sum(p.numel() for p in parameters) == 131072 assert all(p.dtype == torch.float32 for p in parameters) return model, tokenizer, recipe, parameters def score_module(work, lang): if lang == 'cmrc': import nltk # Official mixed-language segmentation calls nltk.word_tokenize. nltk.data.path.insert(0, str(work / 'nltk-data')) path = work / 'sources' / ('squad_evaluate.py' if lang == 'squad' else 'cmrc_evaluate_v6.py') spec = importlib.util.spec_from_file_location('official_' + lang, path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def per_answer(scorer, lang, refs, answer): if lang == 'squad': return max(scorer.compute_exact(r, answer) for r in refs), max(scorer.compute_f1(r, answer) for r in refs) return scorer.calc_em_score(refs, answer), scorer.calc_f1_score(refs, answer) @torch.inference_mode() def generate(model, tokenizer, rows): model.eval() tokenizer.padding_side = 'left' answers = [] torch.cuda.synchronize() started = time.perf_counter() for begin in range(0, len(rows), 8): chunk = rows[begin:begin+8] prompts = [tokenizer.apply_chat_template(row['messages'][:-1], tokenize=False, add_generation_prompt=True, **TEMPLATE) for row in chunk] batch = tokenizer(prompts, add_special_tokens=False, padding=True, return_tensors='pt').to('cuda') result = model.generate(**batch, do_sample=False, max_new_tokens=32, use_cache=True, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id) tokens = result[:, batch['input_ids'].shape[1]:] for row, ids in zip(chunk, tokens): ended = tokenizer.eos_token_id in ids.tolist() answers.append({'id': row['id'], 'prediction': tokenizer.decode(ids, skip_special_tokens=True).strip(), 'references': row['references'], 'reached_limit': not ended and len(ids) == 32}) torch.cuda.synchronize() return answers, time.perf_counter() - started def cleanup(): gc.collect() torch.cuda.empty_cache() def run(work, base_path, lang, method, seed): output = work / lang / f'{method}-{seed}' if (output / 'result.json').exists(): print(f'already complete: {output}', flush=True) return output.mkdir(parents=True, exist_ok=True) model, tokenizer, recipe, params = model_for(method, base_path, work, lang, seed, output) package_identity = {'arti': getattr(arti, '__version__', '3.3.0'), 'torch': torch.__version__, 'transformers': __import__('transformers').__version__, 'peft': __import__('peft').__version__} torch.cuda.reset_peak_memory_stats() train_seconds, loss, cursor = 0., None, None train_trace = [] if method != 'base': torch.cuda.synchronize() started = time.perf_counter() if recipe: # Native recipe execution includes its declared checkpoint/evaluation/export lifecycle. def report(event): if event.get('type') == 'metrics': print(json.dumps({'lang': lang, 'method': method, 'seed': seed, **event}), flush=True) native = recipe.fit(reporter=report) write(output / 'native-receipt.json', native) loss = recipe.session.loss_history[-1] cursor = recipe.train_data.state_dict() train_trace = list(recipe.session.loss_history) else: # Reuse precisely the same caller-owned prepared data and length order. from arti.language_recipe import _prepare_chat_dataset, _training_batches cfg = config(work, lang, seed, output / 'unused') data = _prepare_chat_dataset(cfg, tokenizer) batches = _training_batches(data, cfg) optimizer = torch.optim.AdamW(params, lr=.001, weight_decay=.01, fused=True) model.train() step = 0 while step < 200: for batch in batches: optimizer.zero_grad(set_to_none=True) batch = {k: v.to('cuda') for k, v in batch.items()} objective = model(**batch, use_cache=False).loss objective.backward() torch.nn.utils.clip_grad_norm_(params, 1.) optimizer.step() loss = objective.detach().item() train_trace.append(loss) step += 1 if step % 20 == 0: print(json.dumps({'lang': lang, 'method': method, 'seed': seed, 'step': step, 'loss': loss}), flush=True) if step == 200: break cursor = batches.state_dict() if method == 'peft': model.save_pretrained(output / 'adapter', safe_serialization=True) else: from safetensors.torch import save_file save_file({name: p.detach().cpu().contiguous() for name, p in model.named_parameters() if p.requires_grad}, str(output / 'adapter.safetensors')) torch.cuda.synchronize() train_seconds = time.perf_counter() - started train_peak = {'allocated_bytes': torch.cuda.max_memory_allocated(), 'reserved_bytes': torch.cuda.max_memory_reserved()} torch.cuda.reset_peak_memory_stats() rows = json.loads((work / lang / 'heldout.json').read_text(encoding='utf8')) answers, generation_seconds = generate(model, tokenizer, rows) scorer = score_module(work, lang) for answer in answers: answer['em'], answer['f1'] = per_answer(scorer, lang, answer['references'], answer['prediction']) write(output / 'predictions.json', answers) weight_files = list(output.rglob('*.safetensors')) if method == 'arti': asset_dir = output / 'native-run' / 'model' weight_files = list(asset_dir.rglob('*.st')) + list(asset_dir.rglob('*.safetensors')) result = {'language': lang, 'method': method, 'seed': seed, 'base': BASE, 'revision': REVISION, 'trainable_parameters': sum(p.numel() for p in params), 'steps': 0 if method == 'base' else 200, 'heldout_count': len(answers), 'em': 100 * sum(a['em'] for a in answers) / len(answers), 'f1': 100 * sum(a['f1'] for a in answers) / len(answers), 'generation_limit_count': sum(a['reached_limit'] for a in answers), 'train_wall_seconds': train_seconds, 'train_wall_scope': 'native recipe incl checkpoint/validation/export' if recipe else 'Torch train loop incl scalar loss logging and adapter save', 'heldout_generation_seconds': generation_seconds, 'training_peak': train_peak, 'generation_peak_allocated_bytes': torch.cuda.max_memory_allocated(), 'adapter_payload_bytes': sum(p.stat().st_size for p in weight_files), 'final_train_loss_diagnostic': loss, 'data_cursor': cursor, 'packages': package_identity, 'data_manifest_sha256': digest((work / lang / 'data-manifest.json').read_bytes())} write(output / 'loss-trace.json', train_trace) write(output / 'result.json', result) print(json.dumps({k: result[k] for k in ['language', 'method', 'seed', 'em', 'f1', 'train_wall_seconds']}), flush=True) del model, recipe, params cleanup() def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('operation', choices=['prepare', 'run', 'all']) parser.add_argument('--work', type=Path, default=Path('reading-qa-work')) parser.add_argument('--base-path', type=Path) parser.add_argument('--language', choices=['squad', 'cmrc']) parser.add_argument('--method', choices=METHODS) parser.add_argument('--seed', type=int, default=17) args = parser.parse_args() work = args.work.resolve() work.mkdir(parents=True, exist_ok=True) if args.base_path is None: from huggingface_hub import snapshot_download args.base_path = Path(snapshot_download(BASE, revision=REVISION, cache_dir=work / 'hub-cache')) base_path = args.base_path.resolve() if args.operation in ('prepare', 'all'): prepare(work, base_path) if args.operation == 'run': if not args.language or not args.method: parser.error('run requires --language and --method') run(work, base_path, args.language, args.method, args.seed) elif args.operation == 'all': for lang in ('squad', 'cmrc'): run(work, base_path, lang, 'base', 17) for seed in SEEDS: for method in ('torch', 'peft', 'arti'): run(work, base_path, lang, method, seed) if __name__ == '__main__': main()