File size: 4,608 Bytes
8f825d4 | 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 | """Build an openjev dataset from your own retrieval logs, to fine-tune or re-calibrate Ragzilla.
One JSONL line per question your retriever served, with the passages in the order it returned them
and which of them answered (from clicks, citations, annotators, or a judge you already run):
{"id": "q17", "question": "how long is a cdl permit good for in texas",
"passages": ["...", {"title": "...", "text": "..."}, "..."], "relevant": [0, 3], "group": "optional"}
python my_rows.py my_rows.jsonl --out my_rows # 50 / 10 / 20 / 20 split by group
openjev tod train my_rows --method encoder --model ragzilla --out my_ragzilla \
--loss focal --focal-gamma 2 --focal-alpha 0.5 --brier 0.1 --learning-rate 1e-5 --max-steps 25
openjev tod calibrate my_ragzilla my_rows
openjev tod evaluate my_ragzilla my_rows
Rows are prepared the way Ragzilla's training data was: whitespace collapsed, a titled passage rendered
"<title>\\n<text>", passages over 4,000 characters cut at a word boundary, duplicate passages merged,
empty passages dropped, rows with fewer than 2 passages dropped (and counted). Rows with the same `group`
(a user, a document, the same question) share a split; a small workload needs a calibration split large
enough to choose an operating point on, hence 20% calibration and 20% test by default.
Needs only `pip install openjev-sdk==0.5.0`. The spec below is Ragzilla's own (`spec.json` in this
folder has the same fingerprint), so the rows load into the model as its training rows did.
"""
import argparse
import json
import re
from pathlib import Path
from openjev.tod import Episode, ListSpec, Subset, build_dataset
SPEC = ListSpec('ragzilla', [dict(name='relevant', question='Which passages are relevant to answering the question?',
candidates='passages', version='1')])
MAX_CHARS = 4000
def clean(text):
text = re.sub(r'\s+', ' ', text).strip()
if len(text) > MAX_CHARS:
text = text[:MAX_CHARS].rsplit(' ', 1)[0] + ' ...'
return text
def passage_text(p):
if isinstance(p, dict):
title = re.sub(r'\s+', ' ', p.get('title') or '').strip()
body = clean(p.get('text') or '')
return f'{title}\n{body}' if title and body else body
return clean(str(p))
def decision(row, drops):
"""One SDK decision from a row, or None (counted in `drops`)."""
texts, keep_idx, seen = [], [], {}
for i, p in enumerate(row['passages']):
t = passage_text(p)
if not t:
drops['empty passages'] += 1
continue
if t in seen: # duplicate: the first copy stands for both
if i in set(row.get('relevant', [])):
keep_idx.append(seen[t])
drops['duplicate passages merged'] += 1
continue
seen[t] = len(texts)
if i in set(row.get('relevant', [])):
keep_idx.append(len(texts))
texts.append(t)
if len(texts) < 2:
drops['rows with fewer than 2 passages'] += 1
return None
state = {'question': re.sub(r'\s+', ' ', row['question']).strip(),
'passages': [{'id': f'p{i}', 'text': t} for i, t in enumerate(texts)]}
target = Subset(sorted({f'p{i}' for i in keep_idx}))
try:
return SPEC.decision(str(row['id']), 'relevant', state, target=target, group_id=str(row.get('group') or row['id']))
except ValueError as exc: # e.g. more than 255 passages
drops[f'refused by the SDK: {exc}'[:80]] += 1
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument('rows', type=Path)
ap.add_argument('--out', type=Path, required=True)
ap.add_argument('--fractions', default='0.5,0.1,0.2,0.2', help='train,validation,calibration,test')
ap.add_argument('--seed', default='0')
a = ap.parse_args()
from collections import Counter
drops = Counter()
rows = [json.loads(l) for l in a.rows.open(encoding='utf-8') if l.strip()]
decisions = [d for d in (decision(r, drops) for r in rows) if d is not None]
f = [float(x) for x in a.fractions.split(',')]
fractions = dict(train=f[0], validation=f[1], calibration=f[2], test=f[3])
ds = build_dataset(SPEC, [Episode(d.id, (d,)) for d in decisions], fractions=fractions, seed=a.seed)
ds.write(a.out)
print(json.dumps(dict(out=str(a.out), rows_read=len(rows), decisions=len(decisions), dropped=dict(drops),
spec_fingerprint=SPEC.fingerprint, splits=ds.report['splits']), indent=2))
if __name__ == '__main__':
main()
|