sable-sphynx 0.7.0.dev0__py3-none-any.whl

sphynx/rag_papers/datasets.py

@@ -0,0 +1,291 @@
+"""Pinned native dataset intake and shared-corpus bundles, without inference."""
+from __future__ import annotations
+
+import bz2
+import copy
+import hashlib
+import io
+import json
+import re
+from pathlib import Path
+import zipfile
+
+from .contracts import (MAX_BYTES, canonical, checksum, digest, identity, make_case,
+                        parse, require, validate_case)
+
+BUNDLE_SCHEMA = 'sphynx.rag.paper-dataset.v1'
+
+
+def _read(path, expected_sha256):
+    path = Path(path)
+    require(path.is_file() and not path.is_symlink() and path.stat().st_size <= MAX_BYTES,
+            'Expected a bounded regular source file')
+    raw = path.read_bytes()
+    require(hashlib.sha256(raw).hexdigest() == checksum(expected_sha256), 'Source checksum differs')
+    return raw
+
+
+def _records(raw):
+    if raw[:3] == b'BZh':
+        decompressor = bz2.BZ2Decompressor()
+        raw = decompressor.decompress(raw, max_length=MAX_BYTES+1)
+        require(decompressor.eof and len(raw) <= MAX_BYTES and not decompressor.unused_data,
+                'Compressed source exceeds bounds or contains concatenated streams')
+    try:
+        return parse(raw)
+    except ValueError:
+        return [parse(line) for line in raw.splitlines() if line.strip()]
+
+
+def _document(corpus, text, *, title='', key=None):
+    require(type(text) is str and type(title) is str, 'Document must have exact text/title')
+    key = identity(key or 'doc-'+digest({'text':text,'title':title}))
+    row = dict(id=key, title=title, text=text)
+    require(key not in corpus or corpus[key] == row, 'Conflicting document bytes under one ID')
+    corpus[key] = row
+    return key
+
+
+def _write_bundle(corpus, cases, output_dir, provenance, source=None):
+    output = Path(output_dir)
+    require(not output.exists(), 'Choose a new dataset output directory')
+    output.mkdir(parents=True, exist_ok=False)
+    files = {}
+    def write(name, rows):
+        with (output/name).open('xb') as handle:
+            if isinstance(rows, bytes):
+                handle.write(rows)
+            else:
+                for row in rows:
+                    handle.write(canonical(row)+b'\n')
+        raw=(output/name).read_bytes()
+        require(len(raw)<=MAX_BYTES, 'Dataset file exceeds local byte limit')
+        files[name]=dict(sha256=hashlib.sha256(raw).hexdigest(),bytes=len(raw))
+    write('corpus.jsonl', corpus.values())
+    case_ids=set(); count=0
+    def checked():
+        nonlocal count
+        for case in cases:
+            validate_case(case, corpus)
+            require(case['id'] not in case_ids, 'Duplicate case ID')
+            case_ids.add(case['id']);count+=1
+            yield case
+    write('cases.jsonl', checked())
+    require(count>0, 'Empty benchmark population')
+    if source is not None:
+        write('source.bin', source)
+    manifest=dict(schema_version=BUNDLE_SCHEMA, files=files,
+        counts=dict(documents=len(corpus),cases=count), provenance=provenance,
+        execution=dict(models_run=False,hosted_admitted=False),
+        reference_policy='Native labels are protected annotations, not independent judgments of candidate outputs.')
+    (output/'manifest.json').write_bytes(canonical(manifest)+b'\n')
+    return dict(path=str(output.resolve()), manifest_sha256=hashlib.sha256((output/'manifest.json').read_bytes()).hexdigest(),
+                **manifest['counts'])
+
+
+class PaperDataset:
+    """Verified shared corpus with a streaming, complete case population."""
+    def __init__(self, root, manifest, corpus):
+        self.path,self.manifest,self.corpus=root,manifest,corpus
+        self._identity = digest(dict(manifest=manifest, corpus=corpus))
+        self._manifest_sha256=hashlib.sha256((root/'manifest.json').read_bytes()).hexdigest()
+
+    def cases(self):
+        _read(self.path/'manifest.json',self._manifest_sha256)
+        _read(self.path/'corpus.jsonl',self.manifest['files']['corpus.jsonl']['sha256'])
+        require(digest(dict(manifest=self.manifest, corpus=self.corpus)) == self._identity,
+                'Loaded dataset was mutated; reopen the pinned dataset')
+        # Reverify on each pass, including when callers change a file after load.
+        raw=_read(self.path/'cases.jsonl',self.manifest['files']['cases.jsonl']['sha256'])
+        seen=set()
+        for line in raw.splitlines():
+            if not line.strip():continue
+            case=validate_case(parse(line),self.corpus)
+            require(case['id'] not in seen, 'Duplicate saved case ID')
+            seen.add(case['id']);yield case
+        require(len(seen)==self.manifest['counts']['cases'], 'Saved case census differs')
+
+
+def load_paper_dataset(path, *, expected_manifest_sha256):
+    root=Path(path)
+    manifest=parse(_read(root/'manifest.json',expected_manifest_sha256))
+    require(manifest.get('schema_version')==BUNDLE_SCHEMA, 'Unknown paper dataset schema')
+    require(set(manifest.get('files',{})) in ({'corpus.jsonl','cases.jsonl'},
+                {'corpus.jsonl','cases.jsonl','source.bin'}), 'Unexpected dataset file closure')
+    contents={}
+    for name,receipt in manifest['files'].items():
+        raw=_read(root/name,receipt['sha256'])
+        require(len(raw)==receipt['bytes'], 'Saved byte census differs')
+        if name!='source.bin':contents[name]=raw
+    corpus={}
+    for line in contents['corpus.jsonl'].splitlines():
+        if line.strip():
+            row=parse(line)
+            require(type(row) is dict and set(row)=={'id','title','text'}, 'Saved document fields differ')
+            require(row['id'] not in corpus,'Duplicate saved document ID')
+            _document(corpus,row['text'],title=row['title'],key=row['id'])
+    require(len(corpus)==manifest['counts']['documents'],'Saved corpus census differs')
+    dataset=PaperDataset(root,manifest,corpus)
+    # Consume validation without retaining a second copy of all cases.
+    for _ in dataset.cases():pass
+    return dataset
+
+
+def prepare_canonical_paper_dataset(cases, corpus, output_dir, *, expected_input_sha256, source_id):
+    """Freeze explicit canonical objects, with their corpus and case bytes pinned.
+
+    This records a caller's declared local provenance; it grants no data rights
+    and does not promote model or synthetic annotations into human reviews.
+    """
+    require(type(cases) is list and type(corpus) is dict,'Canonical cases and keyed corpus required')
+    require(digest(dict(cases=cases,corpus=corpus))==checksum(expected_input_sha256),'Canonical input checksum differs')
+    identity(source_id)
+    return _write_bundle(copy.deepcopy(corpus),copy.deepcopy(cases),output_dir,
+        dict(kind='canonical-objects',source_id=source_id,input_sha256=expected_input_sha256))
+
+
+def prepare_official_paper_dataset(source_path, output_dir, *, expected_manifest_sha256):
+    """Bridge the existing complete SciFact/Hotpot bundles without rerunning intake."""
+    from ..official_benchmarks import load_official_benchmark
+    old=load_official_benchmark(source_path,expected_manifest_sha256=expected_manifest_sha256)
+    corpus={d:dict(id=d,title=r.get('title',''),text=r['text']) for d,r in old.corpus.items()}
+    def cases():
+        for qid,query in old.queries.items():
+            native=old.annotations.get(qid,{})
+            refs={'qrels':old.qrels[qid], 'native':copy.deepcopy(native)}
+            if native:
+                refs['answers']=[native['answer']] if type(native.get('answer')) is str else None
+                refs['support_facts']=native.get('supporting_facts') if not native.get('validation_issues') else None
+            yield make_case(qid,query['text'],query.get('context_ids',list(corpus)),references=refs,
+                snapshot=old.manifest['source_sha256'],origin='public',family=old.manifest['dataset_id'],
+                split=old.manifest['upstream']['split'],source_id=old.manifest['dataset_id'])
+    return _write_bundle(corpus,cases(),output_dir,dict(kind='existing-verified-official-bundle',
+        source_manifest_sha256=expected_manifest_sha256,dataset_id=old.manifest['dataset_id']))
+
+
+def prepare_paper_dataset(paper, source_path, output_dir, *, expected_sha256,
+                          source_url=None, split='development', scenario=None):
+    """Read explicit native JSON/JSONL/BZ2 or BEIR ZIP; preserve all native rows.
+
+    A caller pin establishes exact bytes, not independently authenticated origin,
+    license rights or customer authorization. No directory discovery/network.
+    """
+    raw=_read(source_path,expected_sha256)
+    corpus={}; cases=[]
+    def case(cid,question,docs,refs,**extra):
+        cases.append(make_case(identity(str(cid)),question,docs,references=refs,
+            snapshot=expected_sha256,origin='public',family=scenario or paper,split=split,
+            source_id=paper,**extra))
+    if paper=='beir':
+        with zipfile.ZipFile(io.BytesIO(raw)) as archive:
+            require(sum(i.file_size for i in archive.infolist())<=MAX_BYTES,'Expanded BEIR ZIP exceeds bound')
+            roots=[n[:-len('corpus.jsonl')] for n in archive.namelist() if n.endswith('corpus.jsonl')]
+            require(len(roots)==1,'BEIR archive needs exactly one corpus')
+            prefix=roots[0]
+            require(len(archive.namelist())==len(set(archive.namelist())),'Duplicate ZIP member')
+            documents=[parse(x) for x in archive.read(prefix+'corpus.jsonl').splitlines() if x.strip()]
+            queries=[parse(x) for x in archive.read(prefix+'queries.jsonl').splitlines() if x.strip()]
+            for row in documents:_document(corpus,row['text'],title=row.get('title',''),key=row['_id'])
+            qrels={}
+            for line in archive.read(prefix+'qrels/'+split+'.tsv').decode('utf-8').splitlines():
+                if not line.strip() or line.startswith('query-id\t'):continue
+                q,d,g=line.split('\t'); require(d in corpus,'BEIR qrel document absent')
+                require(d not in qrels.setdefault(q,{}),'Duplicate native qrel')
+                qrels[q][d]=int(g)
+            querymap={q['_id']:q for q in queries};require(len(querymap)==len(queries),'Duplicate native query')
+            require(set(qrels)<=set(querymap),'Qrel query absent')
+            for qid,refs in qrels.items():case(qid,querymap[qid]['text'],list(corpus),{'qrels':refs,'native':{'query':querymap[qid]}})
+    else:
+        data=_records(raw)
+        if paper=='qasper':
+            require(type(data) is dict,'QASPER source must map paper IDs')
+            for pid,paperrow in data.items():
+                docs=[]; by_text={}
+                for section in paperrow['full_text']:
+                    for text in section['paragraphs']:
+                        did=_document(corpus,text,title=section.get('section_name',''),key=pid+':'+digest([text,section.get('section_name','')]))
+                        if did not in docs:docs.append(did)
+                        by_text[text]=did
+                for qa in paperrow['qas']:
+                    annotations=[]
+                    for row in qa['answers']:
+                        a=row['answer']
+                        if a['unanswerable']:answer='Unanswerable'; evidence=[]
+                        else:
+                            answer=', '.join(a['extractive_spans']) if a.get('extractive_spans') else a.get('free_form_answer')
+                            if not answer:
+                                require(type(a.get('yes_no')) is bool,'QASPER answer has no recognized form')
+                                answer='Yes' if a['yes_no'] else 'No'
+                            evidence=[e for e in a['evidence'] if 'FLOAT SELECTED' not in e]
+                        annotations.append(dict(answer=answer,evidence_ids=[by_text[e] for e in evidence]
+                            if all(e in by_text for e in evidence) else None))
+                    case(qa['question_id'],qa['question'],docs,{'annotation_sets':annotations,'native':qa},cluster_id=pid)
+        else:
+            if type(data) is dict:data=[data]  # One JSONL row is also a valid JSON object.
+            require(type(data) is list,'Native source must be a JSON array or JSONL records')
+            for index,row in enumerate(data):
+                cid=str(row.get('id',row.get('_id',row.get('interaction_id',index))))
+                if paper=='hotpotqa':
+                    docs=[];titles={}; issues=[]
+                    for title,sentences in row['context']:
+                        did=_document(corpus,''.join(sentences),title=title,key=cid+':'+digest(title))
+                        docs.append(did);titles[title]=(did,len(sentences))
+                    for title,sentence in row['supporting_facts']:
+                        if title not in titles or type(sentence) is not int or not 0<=sentence<titles[title][1]:issues.append([title,sentence])
+                    support=None if issues else row['supporting_facts']
+                    case(cid,row['question'],docs,{'answers':[row['answer']],'support_facts':support,
+                        'support_ids':None if issues else list(dict.fromkeys(titles[t][0] for t,_ in support)),
+                        'native':row})
+                elif paper=='musique':
+                    docs=[];support=[];by_index={}
+                    for p in row['paragraphs']:
+                        did=_document(corpus,p['paragraph_text'],title=p.get('title',''),key=cid+':'+str(p['idx']))
+                        docs.append(did)
+                        by_index[p['idx']]=did
+                        if p['is_supporting']:support.append(did)
+                    answers=list(dict.fromkeys([row['answer']]+row.get('answer_aliases',[])))
+                    hops=row.get('question_decomposition',[])
+                    known=bool(hops) and all(h.get('paragraph_support_idx') in by_index for h in hops)
+                    atom_support={str(h['id']):[by_index[h['paragraph_support_idx']]] for h in hops} if known else None
+                    relations={}
+                    if known:
+                        for hop in hops:
+                            for dependency in re.findall(r'#(\d+)',hop['question']):
+                                require(dependency in atom_support,'MuSiQue dependency points to absent hop')
+                                relations[dependency+'->'+str(hop['id'])]=list(dict.fromkeys(atom_support[dependency]+atom_support[str(hop['id'])]))
+                    case(cid,row['question'],docs,{'answers':answers,'support_ids':support,
+                        'answerable':row['answerable'],'native':row,
+                        'required_atoms':list(atom_support) if known else None,'atom_support':atom_support,
+                        'required_relations':list(relations) if known else None,'relation_support':relations if known else None})
+                elif paper=='alce':
+                    docs=[_document(corpus,d.get('text',d.get('sent','')),title=d.get('title',''),key=cid+':'+str(i+1)) for i,d in enumerate(row['docs'])]
+                    native=copy.deepcopy(row)
+                    if row.get('qa_pairs'):native['asqa_short_answers']=[q['short_answers'] for q in row['qa_pairs']]
+                    if row.get('answers'):native['qampari_answers']=row['answers']
+                    case(cid,row['question'],docs,{'native':native})
+                elif paper=='rgb':
+                    positives=row.get('positive',[]);negatives=row.get('negative',[])
+                    require(type(positives) is list and type(negatives) is list,'RGB context lists required')
+                    positive_groups=positives if positives and type(positives[0]) is list else [positives]
+                    positive_ids=[[_document(corpus,p,key=cid+':p:'+str(g)+':'+str(i)) for i,p in enumerate(group)]
+                                  for g,group in enumerate(positive_groups)]
+                    negative_ids=[_document(corpus,p,key=cid+':n:'+str(i)) for i,p in enumerate(negatives)]
+                    wrong_ids=[_document(corpus,p,key=cid+':w:'+str(i)) for i,p in enumerate(row.get('positive_wrong',[]))]
+                    docs=[d for group in positive_ids for d in group]+negative_ids+wrong_ids
+                    answers=row['answer']; groups=[[answers]] if type(answers) is str else [
+                        a if type(a) is list else [a] for a in answers]
+                    native={**row,'answer_groups':groups,'positive_ids':[d for group in positive_ids for d in group],
+                            'positive_groups':positive_ids,'negative_ids':negative_ids,'wrong_ids':wrong_ids}
+                    case(cid,row.get('query',row.get('question')),docs,{'native':native})
+                elif paper=='crag':
+                    docs=[_document(corpus,d.get('page_result',d.get('page_snippet','')),title=d.get('page_name',''),key=cid+':'+str(i)) for i,d in enumerate(row['search_results'])]
+                    case(cid,row['query'],docs,{'answers':[row['answer']],'native':row},context={'as_of':row.get('query_time','')})
+                elif paper=='lost_in_the_middle':
+                    docs=[_document(corpus,d['text'],title=d.get('title',''),key=cid+':'+str(i)) for i,d in enumerate(row['ctxs'])]
+                    support=[did for did,d in zip(docs,row['ctxs']) if d.get('hasanswer') is True]
+                    case(cid,row['question'],docs,{'answers':row['answers'],'support_ids':support,'native':row})
+                else:
+                    require(False,'No native loader for this paper; use canonical case/judgment inputs')
+    provenance=dict(paper=paper,source_sha256=expected_sha256,source_url=source_url,split=split,
+                    origin_basis='caller-pinned native bytes; preserve independent publisher/license receipt')
+    return _write_bundle(corpus,cases,output_dir,provenance,source=raw)

sphynx/rag_papers/experiments.py

@@ -0,0 +1,116 @@
+"""Frozen controlled interventions, separate from ordinary benchmark grading."""
+from __future__ import annotations
+import copy
+import math
+import random
+
+from .contracts import candidate_input, digest, require, validate_case
+from .datasets import _write_bundle
+
+
+def plan_lost_in_the_middle(case, corpus, *, evidence_id, positions, seed,
+                            render_prompt, count_tokens, tokenizer_revision, max_tokens):
+    """Move one evidence document through the same frozen distractor population.
+
+    The caller supplies its actual prompt renderer and tokenizer. Every prompt
+    must fit the same declared budget; no text is truncated or resampled. Exact
+    text hashes, token counts and randomized execution order are retained.
+    """
+    validate_case(case,corpus)
+    require(evidence_id in case['document_ids'] and type(positions) is list and positions
+            and len(set(positions))==len(positions) and all(type(p) is int and 0<=p<len(case['document_ids']) for p in positions),
+            'Invalid evidence ID or zero-based positions')
+    require(type(seed) is int and callable(render_prompt) and callable(count_tokens)
+            and type(tokenizer_revision) is str and tokenizer_revision
+            and type(max_tokens) is int and max_tokens>0, 'Explicit prompt/tokenizer and budget required')
+    others=[d for d in case['document_ids'] if d!=evidence_id]
+    random.Random(seed).shuffle(others)
+    trials=[]
+    for position in positions:
+        c=copy.deepcopy(case); order=list(others);order.insert(position,evidence_id)
+        c['id']='litm-'+digest([case['id'],position,seed])[:32]; c['document_ids']=order
+        c['metadata']['scenario_family']='lost_in_the_middle.position'
+        prompt=render_prompt(candidate_input(c,corpus))
+        require(type(prompt) is str,'Prompt renderer must return exact text')
+        tokens=count_tokens(prompt)
+        require(type(tokens) is int and 0<tokens<=max_tokens,'Position trial exceeds token budget')
+        trials.append(dict(case=validate_case(c,corpus),position=position,prompt_sha256=digest(prompt),tokens=tokens))
+    random.Random(seed+1).shuffle(trials)
+    return dict(schema_version='sphynx.rag.position-plan.v1', method='lost_in_the_middle',
+        source_case_sha256=digest(case), corpus_sha256=digest({d:corpus[d] for d in case['document_ids']}),
+        evidence_id=evidence_id,seed=seed,tokenizer_revision=tokenizer_revision,max_tokens=max_tokens,
+        fixed_distractor_order=others,trials=trials,models_run=False)
+
+
+def plan_rgb(case, corpus, *, testbed, passage_count=5, noise_rate=0, correct_rate=0, seed=0):
+    """Freeze RGB testbed selections and shuffle seed, failing on missing inputs.
+
+    Selection conventions are explicit Sphynx variants. The original scripts
+    also contain fallback sampling behavior; this planner refuses to silently
+    change a requested experimental condition when a pool is too small.
+    """
+    validate_case(case,corpus)
+    require(testbed in ('noise','rejection','integration','counterfactual'), 'Unknown RGB testbed')
+    require(type(passage_count) is int and passage_count>0 and type(seed) is int,'Invalid RGB trial size/seed')
+    for rate in (noise_rate,correct_rate):require(type(rate) in (int,float) and math.isfinite(rate) and 0<=rate<=1,'Invalid RGB rate')
+    n=case['references']['native']; rng=random.Random(seed)
+    positives=n.get('positive_ids',[]); negatives=n.get('negative_ids',[])
+    negative_count=passage_count if testbed=='rejection' else math.ceil(passage_count*noise_rate)
+    require(negative_count<=len(negatives),'Insufficient frozen negative documents')
+    selected=list(negatives[:negative_count]); positive_count=passage_count-negative_count
+    if testbed=='integration':
+        groups=copy.deepcopy(n.get('positive_groups',[]))
+        require(groups and all(groups) and positive_count>=len(groups),'Integration needs evidence from every group')
+        for group in groups:rng.shuffle(group)
+        round_robin=[g[i] for i in range(max(map(len,groups))) for g in groups if i<len(g)]
+        require(positive_count<=len(round_robin),'Insufficient integration evidence')
+        selected+=round_robin[:positive_count]
+    elif testbed=='counterfactual':
+        correct_count=math.ceil(passage_count*correct_rate); wrong_count=positive_count-correct_count
+        wrong=n.get('wrong_ids',[])
+        require(0<=wrong_count and len(wrong)==len(positives) and len(wrong)>=positive_count,
+                'Counterfactual evidence must align with native positive documents')
+        indices=rng.sample(range(len(positives)),positive_count)
+        selected += [wrong[i] for i in indices[:wrong_count]]+[positives[i] for i in indices[wrong_count:]]
+    else:
+        require(positive_count<=len(positives),'Insufficient frozen positive documents')
+        selected+=positives[:positive_count]
+    require(len(set(selected))==passage_count and set(selected)<=set(case['document_ids']),'RGB selection is incomplete')
+    rng.shuffle(selected); c=copy.deepcopy(case)
+    c['id']='rgb-'+digest([case['id'],testbed,selected,seed])[:32];c['document_ids']=selected
+    c['metadata']['scenario_family']='rgb.'+testbed
+    return dict(schema_version='sphynx.rag.rgb-plan.v1', method='rgb', source_case_sha256=digest(case),
+        testbed=testbed,passage_count=passage_count,noise_rate=1 if testbed=='rejection' else noise_rate,
+        correct_rate=correct_rate,seed=seed,case=validate_case(c,corpus),models_run=False)
+
+
+def prepare_intervention_dataset(plan, corpus, output_dir):
+    """Materialize a plan for capture_paper_candidate followed by run_paper_suite."""
+    require(plan.get('schema_version') in ('sphynx.rag.position-plan.v1','sphynx.rag.rgb-plan.v1','sphynx.rag.musique-plan.v1'),'Unknown intervention plan')
+    cases=[t['case'] for t in plan['trials']] if 'trials' in plan else [plan['case']]
+    return _write_bundle(corpus,cases,output_dir,dict(kind='controlled-intervention',plan=copy.deepcopy(plan),plan_sha256=digest(plan)))
+
+
+def plan_musique(case, corpus, *, seed=0):
+    """Execute full-context and leave-one-hop-evidence-out connected-QA trials.
+
+    Removed annotated evidence does not prove remaining context insufficient:
+    redundant support and model prior knowledge need separate interpretation.
+    Source answerability and original gold stay intact across interventions.
+    """
+    validate_case(case,corpus)
+    support=case['references']['atom_support'];required=case['references']['required_atoms']
+    require(type(seed) is int and required and support and all(a in support and support[a] for a in required),
+            'Frozen required-hop evidence annotations are needed')
+    trials=[]
+    for atom in [None]+required:
+        c=copy.deepcopy(case)
+        removed=[] if atom is None else support[atom]
+        c['document_ids']=[d for d in case['document_ids'] if d not in removed]
+        c['id']='musique-'+digest([case['id'],atom,seed])[:32]
+        c['metadata']['scenario_family']='musique.connected.full' if atom is None else 'musique.connected.ablation'
+        if atom is not None:c['references']['expected_action']=None
+        trials.append(dict(case=validate_case(c,corpus),removed_hop=atom,removed_ids=removed))
+    random.Random(seed).shuffle(trials)
+    return dict(schema_version='sphynx.rag.musique-plan.v1',method='musique',source_case_sha256=digest(case),
+        seed=seed,trials=trials,models_run=False,insufficiency_inferred_from_removal=False)