Spaces:
Sleeping
Sleeping
File size: 7,381 Bytes
c2b7494 | 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Local MLX reviewer for Apple silicon; no API calls or model substitution."""
from functools import lru_cache
import hashlib
import json
import os
from pathlib import Path
import re
from time import perf_counter
from uuid import uuid4
def visible_response(text):
"""Exclude reasoning, including a thinking block preopened by the template."""
text = re.sub(r'<think>.*?</think>', '', text, flags=re.S)
if '</think>' in text:
text = text.split('</think>',1)[1]
if '<think>' in text:
text = text.split('<think>',1)[0]
return text.strip()
def first_response(text):
"""Read a complete JSON envelope, allowing subsequent prose or repeated output.
Never repair incomplete JSON or extract a nested classification from a broken
envelope. Generation and classification settings remain unchanged.
"""
text=visible_response(text)
start=text.find('{')
if start<0:raise ValueError('No complete JSON envelope in local response')
result,_=json.JSONDecoder().raw_decode(text[start:])
if not isinstance(result,dict):raise ValueError('Local response must be a JSON object')
if not isinstance(result.get('classifications'),list):
raise ValueError('Local response must contain a classifications array')
return result
def candidate_parts(user_prompt):
marker='Classify each of the following candidate codes:\n'
if marker not in user_prompt:return None
prefix,body=user_prompt.split(marker,1)
lines,suffix=body.split('\n\nReturn a JSON object',1)
candidates=lines.splitlines()
if not candidates or any(not line.startswith('- code: ') for line in candidates):return None
return prefix+marker,candidates,'\n\nReturn a JSON object'+suffix
def validate_candidates(response,user_prompt):
parts=candidate_parts(user_prompt)
if parts is None:return
expected=[]
for line in parts[1]:
code,vocabulary,_=line.removeprefix('- code: ').split(' | ',2)
expected.append((code,vocabulary.removeprefix('vocabulary: ')))
actual=[(r.get('code'),r.get('vocabulary')) for r in response['classifications']]
if len(actual)!=len(expected) or set(actual)!=set(expected):
raise ValueError('Local answer does not cover every requested code exactly once')
@lru_cache(maxsize=1)
def load_mlx_reviewer(model_name):
from mlx_vlm import load
from src.llm.model_paths import prepared_model_path
return load(prepared_model_path(model_name))
def call_mlx_json(system_prompt, user_prompt, model_name, max_tokens=8192, reasoning_effort='medium',
thinking_budget=1024):
from mlx_vlm import stream_generate
model, processor = load_mlx_reviewer(model_name)
messages = [{'role':'system','content':system_prompt},{'role':'user','content':user_prompt}]
prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
enable_thinking=reasoning_effort not in (None,'none'))
folder=Path(os.environ.get('LOCAL_RESPONSE_DIR',Path(__file__).resolve().parents[2]/'results/local_responses'))
folder.mkdir(parents=True,exist_ok=True)
prompt_hash=hashlib.sha256(prompt.encode()).hexdigest()
failure=folder/'failed'/f'{prompt_hash}.json'
replay=json.loads(failure.read_text()) if failure.exists() else None
if replay and (replay.get('model'),replay.get('max_tokens'),replay.get('thinking_budget'))!=(model_name,max_tokens,thinking_budget):
replay=None
if replay:
print('[Qwen reviewer] Reusing audited failed output; subdividing its candidate batch.',flush=True)
chunks = []; started = perf_counter(); reported = 0; generated_tokens=0
finish_reason='model_stop_or_limit'
outputs=() if replay else stream_generate(model, processor, prompt=prompt, max_tokens=max_tokens,
temperature=.6, top_p=.95, top_k=20, seed=42,
enable_thinking=reasoning_effort not in (None,'none'),
thinking_budget=thinking_budget)
for output in outputs:
chunks.append(output.text)
generated_tokens=output.generation_tokens
if output.generation_tokens >= reported+256:
reported = output.generation_tokens
print(f'[Qwen reviewer] {reported} tokens generated, {perf_counter()-started:.1f}s',flush=True)
current=''.join(chunks)
if '</think>' in current or reasoning_effort in (None,'none'):
try:
parsed=first_response(current)
validate_candidates(parsed,user_prompt)
except json.JSONDecodeError as exc:
# An unescaped control character inside a JSON string cannot
# become valid by generating more tokens. Other incomplete
# prefixes remain eligible to finish normally.
if exc.msg.startswith('Invalid control character'):
finish_reason='invalid_json_syntax'
if hasattr(outputs,'close'):outputs.close()
break
except (ValueError,TypeError):pass
else:
finish_reason='complete_json'
if hasattr(outputs,'close'):outputs.close()
break
text = replay['raw_response'] if replay else ''.join(chunks)
path=folder/(uuid4().hex+'.json')
audit={'model':model_name,'raw_response':text,'generation_tokens':generated_tokens,
'max_tokens':max_tokens,'thinking_budget':thinking_budget,'seconds':perf_counter()-started,
'prompt_sha256':prompt_hash,'replayed_failed_response':bool(replay),'finish_reason':finish_reason,
'implementation_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest()}
try:
result=first_response(text)
validate_candidates(result,user_prompt)
audit['status']='parsed'
return result
except (ValueError,TypeError) as exc:
audit.update(status='parse_error',error=str(exc))
failure.parent.mkdir(exist_ok=True)
failure.write_text(json.dumps(audit,ensure_ascii=False,indent=2))
print(f'[Qwen reviewer] Raw failed response saved: {path}',flush=True)
parts=candidate_parts(user_prompt)
if parts and len(parts[1])>1:
prefix,lines,suffix=parts;middle=len(lines)//2
audit['recovery']='subdivide_batch'
# Persist the failure before recovery, including if recovery is interrupted.
path.write_text(json.dumps(audit,ensure_ascii=False,indent=2))
print(f'[Qwen reviewer] Retrying failed batch as {middle} + {len(lines)-middle} candidates.',flush=True)
recovered=[]
for subset in (lines[:middle],lines[middle:]):
answer=call_mlx_json(system_prompt,prefix+'\n'.join(subset)+suffix,model_name,
max_tokens,reasoning_effort,thinking_budget)
recovered.extend(answer['classifications'])
result={'classifications':recovered}
validate_candidates(result,user_prompt)
audit['status']='recovered_by_subdivision'
return result
raise
finally:
path.write_text(json.dumps(audit,ensure_ascii=False,indent=2))
|