Harness step 3b: LLM-labelled segmentation on STT timestamps ("STT cuts, LLM labels").
"""Harness step 3b: LLM-labelled segmentation on STT timestamps ("STT cuts, LLM labels").
The diarised Soniox transcript is rendered as numbered turns with start/end seconds; a TEXT LLM assigns each
subtask a first/last turn. Windows therefore inherit Soniox's timestamps. Compared against the tablet's
per-subtask start/end (relative durations + one fitted offset).
Run: uv run --with requests --with pandas python harness/03b_segmenter_llm.py [model] [max_files]
"""
import os, sys, json, glob, re, time, requests, concurrent.futures as cf
import pandas as pd
ROOT = "/Users/haroonyasin/Documents/Projects/April 2026/Rumi 10 April 2026"
KEY = [l.split('=', 1)[1].strip() for l in open(os.path.join(ROOT, '.env')) if l.startswith('NIETE_STAGING_OPENROUTER_API_KEY=')][-1]
MODEL = sys.argv[1] if len(sys.argv) > 1 else 'google/gemini-2.5-flash'
MAXF = int(sys.argv[2]) if len(sys.argv) > 2 else 10000
os.makedirs('harness/segments_llm', exist_ok=True)
SUBS = ['listcomp_urd', 'letterid_urd', 'idwrd_urd', 'orf_urd', 'rdcomp_urd', 'listcomp_eng', 'letterid_eng', 'pw_eng', 'idwrd_eng', 'orf_eng', 'rdcomp_eng',
'idnummag', 'numrep', 'blfluency', 'bleffcomp', 'wrdpblm', 'patterns']
PROMPT = """You are segmenting a transcript of an enumerator administering an EGRA/EGMA battery to one child in a Pakistani school. The transcript is diarised and numbered; each turn has [start-end seconds] and a speaker id. The battery order is usually:
URDU: listcomp_urd (enumerator reads a story twice, asks 6 questions) → letterid_urd (child names Urdu letters for 60 s after "start") → idwrd_urd (child reads a list of Urdu words) → orf_urd (child reads an Urdu story/passage aloud) → rdcomp_urd (questions about that passage)
ENGLISH: listcomp_eng (English story + questions) → letterid_eng (English letters, 60 s) → pw_eng (made-up/nonsense words) → idwrd_eng (English word list) → orf_eng (English passage) → rdcomp_eng (questions)
MATH: idnummag (which number is bigger / identify numbers) → numrep (represent numbers, sticks/bundles, place value) → blfluency (rapid addition then subtraction, timed) → bleffcomp (computation like 25+7) → wrdpblm (word problems) → patterns (shape/number patterns)
Some subtasks may be skipped. For every subtask you can identify, output the turn number where the CHILD'S RESPONSE begins (for timed reads: the child's first letter/word after the enumerator says start; for question tasks: the first question) and the turn number where it ends. Also output the turn number of the enumerator's instruction that introduced it.
Return ONLY JSON: {"segments":[{"sub":"letterid_urd","instr_turn":12,"first_turn":13,"last_turn":14,"confidence":0.9}, ...]}"""
def render(path):
tr = json.load(open(path))
toks = tr.get('tokens', [])
turns = []
for t in toks:
s, e, sp, tx = t.get('start_ms', 0) / 1000, t.get('end_ms', 0) / 1000, t.get('speaker'), t.get('text', '')
if turns and turns[-1]['spk'] == sp and s - turns[-1]['e'] < 1.5:
turns[-1]['text'] += tx; turns[-1]['e'] = e
else:
turns.append({'spk': sp, 's': s, 'e': e, 'text': tx})
lines = [f"{i} [{tu['s']:.0f}-{tu['e']:.0f}] spk{tu['spk']}: {tu['text'].strip()[:400]}" for i, tu in enumerate(turns)]
return turns, '\n'.join(lines)
def label(path):
out = f"harness/segments_llm/{MODEL.replace('/', '_')}__{os.path.basename(path)}"
if os.path.exists(out):
return json.load(open(out))
turns, text = render(path)
for attempt in range(3):
try:
r = requests.post('https://openrouter.ai/api/v1/chat/completions', headers={'Authorization': f'Bearer {KEY}'},
json={'model': MODEL, 'messages': [{'role': 'user', 'content': PROMPT + '\n\nTRANSCRIPT:\n' + text}], 'max_tokens': 16000, 'temperature': 0}, timeout=300)
j = r.json(); c = j['choices'][0]['message']['content']
m = re.search(r'\{.*\}', c, re.S); segs = json.loads(m.group(0))['segments']
res = {'file': os.path.basename(path), 'model': MODEL, 'cost': j.get('usage', {}).get('cost'), 'segments': []}
for sg in segs:
try:
ft, lt, it = int(sg['first_turn']), int(sg['last_turn']), int(sg.get('instr_turn', sg['first_turn']))
res['segments'].append({'sub': sg['sub'], 'instr_s': turns[it]['s'], 'start_s': turns[ft]['s'], 'end_s': turns[lt]['e'], 'confidence': sg.get('confidence')})
except Exception:
pass
json.dump(res, open(out, 'w'), ensure_ascii=False)
return res
except Exception as ex:
time.sleep(2 + attempt * 3)
return None
def compare(res, gt):
st = gt['subtask_times']
T = lambda k, w: pd.to_datetime(st[k][w]) if st.get(k, {}).get(w) else None
tab = {}
for k in SUBS:
a, b = T(k, 'start'), T(k, 'end')
if a is not None and b is not None and 0 < (b - a).total_seconds() < 900:
tab[k] = (a, (b - a).total_seconds())
got = {s['sub']: s for s in res['segments']}
common = [k for k in tab if k in got]
if len(common) < 2:
return {'n_tablet': len(tab), 'n_llm': len(got), 'n_common': len(common)}
# fit offset: median over common of (llm start - tablet start seconds-since-form-start)
fs = pd.to_datetime(gt['form_start'])
diffs = [got[k]['start_s'] - (tab[k][0] - fs).total_seconds() for k in common]
off = sorted(diffs)[len(diffs) // 2]
errs = {k: round(got[k]['start_s'] - ((tab[k][0] - fs).total_seconds() + off), 1) for k in common}
lens = {k: (round(got[k]['end_s'] - got[k]['start_s'], 1), round(tab[k][1], 1)) for k in common}
within5 = sum(1 for v in errs.values() if abs(v) <= 5); within15 = sum(1 for v in errs.values() if abs(v) <= 15)
timed = [k for k in common if k in ('letterid_urd', 'letterid_eng', 'idwrd_urd', 'idwrd_eng', 'pw_eng', 'orf_urd', 'orf_eng', 'blfluency')]
len_err = [abs(lens[k][0] - lens[k][1]) for k in timed]
return {'n_tablet': len(tab), 'n_llm': len(got), 'n_common': len(common), 'within5': within5, 'within15': within15,
'mean_abs_start_err': round(sum(abs(v) for v in errs.values()) / len(errs), 1), 'timed_len_err': round(sum(len_err) / len(len_err), 1) if len_err else None,
'errs': errs, 'lens': lens, 'cost': res.get('cost')}
if __name__ == '__main__':
files = sorted(glob.glob('harness/soniox/*.json'))[:MAXF]
files = [f for f in files if os.path.exists('harness/gt/' + os.path.basename(f))]
print(MODEL, len(files), 'files', flush=True)
rows = []
with cf.ThreadPoolExecutor(4) as ex:
for res in ex.map(label, files):
if not res:
continue
gt = json.load(open('harness/gt/' + res['file']))
c = compare(res, gt); c['file'] = res['file'][:20]; rows.append(c)
print(res['file'][:20], {k: v for k, v in c.items() if k in ('n_tablet', 'n_llm', 'n_common', 'within5', 'within15', 'mean_abs_start_err', 'timed_len_err', 'cost')}, flush=True)
df = pd.DataFrame(rows)
df.to_csv(f"harness/segmenter_llm_{MODEL.replace('/', '_')}.csv", index=False)
if 'within5' in df:
d = df.dropna(subset=['within5'])
print(f"\nSUMMARY {MODEL}: files={len(d)} subtasks tablet={d.n_tablet.sum()} llm={d.n_llm.sum()} common={d.n_common.sum()} "
f"start within 5s={d.within5.sum()} ({100 * d.within5.sum() / d.n_common.sum():.0f}%) within 15s={d.within15.sum()} ({100 * d.within15.sum() / d.n_common.sum():.0f}%) "
f"mean|start err|={d.mean_abs_start_err.mean():.1f}s timed-len err={d.timed_len_err.mean():.1f}s cost/file=${d.cost.mean():.4f}")