File size: 4,559 Bytes
7bcd331 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env python3
"""Check native inference against a fixed subset of saved benchmark outputs."""
import argparse
from collections import defaultdict
import importlib.metadata
import json
import platform
from pathlib import Path
import sys
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from inference import ChatModel
from inference_native.chat_format import serialize
from inference_native.tokenization import FIXTURES, SPECIALS
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--weights', type=Path)
parser.add_argument('--output', required=True, type=Path)
args = parser.parse_args()
model = ChatModel(weights=args.weights)
guards = []
for text in FIXTURES:
assert model.tokenizer.decode(model.tokenizer.encode(text)) == text
assert not set(model.tokenizer.encode(text)) & {0, 1, 2, 3}
guards.append('Tokenizer round trips and literal structural-token isolation')
messages = [{'role': 'user', 'content': 'hello'}]
expected = [2]+model.tokenizer.encode('user\n')+model.tokenizer.encode('hello')+[3, 2]+model.tokenizer.encode('assistant\n')
assert model.prepare(messages, 128).tolist() == expected
guards.append('Explicit chat token boundaries')
for invalid in [[], [{'role':'assistant','content':'hi'}], [{'role':'user','content':''}],
[{'role':'user','content':'a'},{'role':'user','content':'b'}],
[{'role':'system','content':'s'}], [{'role':'tool','content':'x'}]]:
try: model.prepare(invalid, 128)
except ValueError: pass
else: raise AssertionError(f'Accepted invalid conversation: {invalid}')
guards.append('Invalid roles, empty messages and turn orders rejected')
for invalid in [0, 2048, True, -1]:
try: model.prepare(messages, invalid)
except ValueError: pass
else: raise AssertionError('Accepted invalid output budget')
try: model.prepare([{'role':'user','content':'word '*2200}], 128)
except ValueError: pass
else: raise AssertionError('Context overflow was not rejected')
guards.append('Output budget bounds and context overflow rejected')
for value in [float('nan'), float('inf'), -1, 2.1]:
try: model.generate(messages, temperature=value)
except ValueError: pass
else: raise AssertionError('Accepted invalid temperature')
guards.append('Non-finite and out-of-range temperatures rejected')
groups = defaultdict(list)
for line in (ROOT/'results/public/chat/responses.jsonl').read_text().splitlines():
row = json.loads(line)
groups[row['case']['category']].append(row)
# Chosen by fixed positions, including failures; not by answer correctness.
selected = [rows[i] for rows in groups.values() for i in (0, 1, 5, 15) if i < len(rows)]
comparisons = []
for row in selected:
case = row['case']; prior = row['response']
result = model.generate(case['messages'], max_new_tokens=case['max_tokens'])
matches = result['text'] == prior['text'] and result['stop_reason'] == prior['stop_reason'] and result['output_tokens'] == prior['tokens']
comparisons.append({'id':case['id'], 'matches_saved_fp32':matches, 'previous':prior, 'observed':result})
print(f"{case['id']}: {'match' if matches else 'DIFFERENT'}", flush=True)
# Verify a sampled invocation is reproducible, without asserting its quality.
sample = model.generate(messages, max_new_tokens=32, temperature=.6, seed=2026)
repeated = model.generate(messages, max_new_tokens=32, temperature=.6, seed=2026)
assert sample['token_ids'] == repeated['token_ids']
guards.append('Fixed-seed sampling is repeatable in this runtime')
report = {'passed': all(r['matches_saved_fp32'] for r in comparisons),
'full_benchmark_rerun':False, 'weight_sha256':model.config['weight_sha256'],
'environment': {'python':platform.python_version(), 'macos':platform.mac_ver()[0],
'machine':platform.machine(), **{k:importlib.metadata.version(k) for k in ['mlx','numpy','tokenizers']}},
'checks':guards, 'saved_response_comparisons':len(comparisons),
'matching_responses':sum(r['matches_saved_fp32'] for r in comparisons),
'comparisons':comparisons}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(report, indent=2)+'\n')
raise SystemExit(not report['passed'])
if __name__ == '__main__':
main()
|