ChristophSchuhmann's picture
Release merged NotebookLM full FT + rank-64 final 652
9771537 verified
Raw History Blame Contribute Delete
3.57 kB
"""Random access to 32-book continuous timelines; model consumes first twelve."""
from collections import OrderedDict
import hashlib,io,json,sqlite3
import numpy as np
from common import *
class Data:
def __init__(self,pilot=False,path=None):
if pilot:raise RuntimeError('No partial-corpus production mode')
self.root=RUN/'large';self.ready=json.loads((self.root/'DATA_READY.json').read_text())
self.db=sqlite3.connect('file:'+str(path or self.root/'mix.sqlite')+'?mode=ro&immutable=1',uri=True)
self.db.row_factory=sqlite3.Row;self.files=OrderedDict();self.cache=OrderedDict()
self.count=self.db.execute('SELECT count(*) FROM presentations').fetchone()[0]
def episode(self,uid):
if uid not in self.cache:
r=self.db.execute('SELECT * FROM episodes WHERE uid=?',(uid,)).fetchone()
if not r:raise KeyError(uid)
if r['tar'] not in self.files:
self.files[r['tar']]=open(r['tar'],'rb')
if len(self.files)>16:self.files.popitem(last=False)[1].close()
f=self.files[r['tar']];f.seek(r['code_offset']);blob=f.read(r['code_size'])
if hashlib.sha256(blob).hexdigest()!=r['code_sha']:raise RuntimeError('Continuous codec timeline hash changed')
codes=np.load(io.BytesIO(blob),allow_pickle=False)
if codes.ndim!=2 or codes.shape[1]!=32 or codes.dtype.kind not in 'iu' or codes.min()<0 or codes.max()>1023:raise ValueError('Invalid podcast codec array')
self.cache[uid]=codes
if len(self.cache)>8:self.cache.popitem(last=False)
self.cache.move_to_end(uid);return self.cache[uid]
def row(self,key,condition='reference'):
r=self.db.execute('SELECT * FROM samples WHERE key=?',(int(key),)).fetchone()
if not r:raise KeyError(key)
ex=json.loads(r['example_json']);codes=self.episode(r['episode'])
a,b=ex['target_frame_start'],ex['target_frame_end'];ca,cb=ex['context_frame_start'],ex['context_frame_end']
if not 0<=ca<=cb<=a<b<=len(codes) or cb-ca>750:raise ValueError('Future context or invalid boundaries')
refs={}
for speaker,q in ex['speaker_references'].items():
ra,rb=q['start_frame'],q['end_frame']
if not 0<=ra<rb<=a or not 63<=rb-ra<=187:raise ValueError('Reference must be past-only and 5–15s')
refs[speaker]=codes[ra:rb,:12].copy()
if set(refs)!={'S1','S2'}:raise ValueError('Two past identity references required')
description='\n'.join(f"{q['speaker_id']}: {max(0,q['start_s']-ca/12.5):.2f}–{q['end_s']-ca/12.5:.2f} seconds in previous audio" for q in ex['context_speaker_intervals'])
return dict(uid=r['uid'],episode=r['episode'],source=r['source'],language=r['language'],plain=r['plain'],script=r['inline'],
instruction='GENERAL: '+r['general']+'\nSCRIPT:\n'+r['inline'],frames=r['frames'],target_codes=codes[a:b,:12].copy(),
speaker_references=refs,reference_codes=None,context_codes=codes[ca:cb,:12].copy(),context_description=description,
conditioning='two_speaker_references_and_past_context',target_turns=ex['target_turns'])
def presentation(self,pos):
r=self.db.execute('SELECT * FROM presentations WHERE pos=?',(int(pos),)).fetchone()
if not r:raise KeyError(pos)
key=r['sample_key']
if key is None:key=self.ready['canary_keys'][0]
return self.row(key),r['sample_key'] is not None
def valid_count(self,start,width):
return self.db.execute('SELECT count(*) FROM presentations WHERE pos>=? AND pos<? AND sample_key IS NOT NULL',(start,start+width)).fetchone()[0]
def canaries(self):return self.ready['canary_keys']
def validation_keys(self):return self.ready['validation_keys']
def close(self):
for f in self.files.values():f.close()
self.db.close()