82 lines
4.4 KiB
Python
82 lines
4.4 KiB
Python
"""Independent FTS5/BM25 and E5 candidates over witnessed source assertions.
|
|
|
|
The checkpoint is an evidence store, not an alternative to text retrieval.
|
|
Named finite scopes are expanded after ranking so a bounded candidate list
|
|
cannot silently remove an assignment goal or one side of a comparison.
|
|
"""
|
|
from functools import lru_cache
|
|
import base64
|
|
import sqlite3
|
|
import threading
|
|
import zlib
|
|
|
|
from tools.answers.context_policy import normalized
|
|
|
|
_lock = threading.RLock()
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def lexical_index(version, texts):
|
|
connection = sqlite3.connect(':memory:', check_same_thread=False)
|
|
connection.execute('CREATE VIRTUAL TABLE evidence USING fts5(id UNINDEXED, text, tokenize="unicode61 remove_diacritics 2")')
|
|
connection.executemany('INSERT INTO evidence VALUES (?,?)',texts)
|
|
return connection
|
|
|
|
|
|
def retrieve_units(question, index, spec, *, limit=5):
|
|
from tools.knowledge.query import _terms
|
|
from tools.knowledge.evidence import claim_eligible
|
|
people={p.name for p in spec.people if p.role=='student'}
|
|
similar=spec.intent=='similar'
|
|
def in_scope(row):
|
|
if not spec.matches_work(row['person'],row.get('work_label'),row.get('work_type'),row.get('work_year')):return False
|
|
if spec.period and spec.period.field!='work_year':
|
|
value=row.get(spec.period.field)
|
|
if not value or not spec.period.start<=value<=spec.period.end:return False
|
|
return similar or not people or row['person'] in people
|
|
candidates=[r for r in index['units'] if in_scope(r)]
|
|
eligible={r['id'] for r in candidates}
|
|
texts=tuple((r['id'],r['text']+' '+' '.join(r.get('concepts',[]))) for r in index['units'])
|
|
connection=lexical_index(index['version'],texts)
|
|
terms=_terms(question)
|
|
if similar:
|
|
# The reference's actual substantive topics seed the peer search.
|
|
from tools.knowledge.concepts import substantive_concept
|
|
terms={normalized(c) for r in candidates if r['person'] in people and claim_eligible(r)
|
|
for c in r.get('concepts',[]) if substantive_concept(c)} or terms
|
|
expression=' OR '.join('"'+term.replace('"','""')+'"*' for term in sorted(terms))
|
|
with _lock:
|
|
lexical=[r[0] for r in connection.execute(
|
|
'SELECT id FROM evidence WHERE evidence MATCH ? ORDER BY bm25(evidence),id',
|
|
(expression,)) if r[0] in eligible] if expression else []
|
|
semantic={}
|
|
if index.get('vectors'):
|
|
import numpy as np
|
|
from tools.rag.embedding_utils import embed_query,model_profile,embedding_model_name
|
|
if model_profile(embedding_model_name())['fingerprint']!=index['embedding_fingerprint']:
|
|
raise RuntimeError('Content embeddings belong to another model')
|
|
matrix=np.frombuffer(zlib.decompress(base64.b64decode(index['vectors'])),dtype='<f4').reshape(index['vector_shape'])
|
|
if len(matrix)!=len(index['units']):raise ValueError('Invalid content vector count')
|
|
anchors=[i for i,r in enumerate(index['units']) if r['person'] in people and in_scope(r) and claim_eligible(r)] if similar else []
|
|
vector=np.mean(matrix[anchors],axis=0) if anchors else embed_query(question)
|
|
norm=float(np.linalg.norm(vector))
|
|
if not norm:raise ValueError('Invalid query/reference embedding')
|
|
scores=matrix @ (vector/norm)
|
|
semantic={r['id']:float(s) for r,s in zip(index['units'],scores) if r['id'] in eligible}
|
|
vector_rank=sorted(semantic,key=lambda key:(-semantic[key],key))
|
|
budget=min(256,max(80,limit*16))
|
|
fused={}
|
|
for ranking in (lexical[:budget],vector_rank[:budget]):
|
|
for position,key in enumerate(ranking,1):fused[key]=fused.get(key,0)+1/(60+position)
|
|
ranked=sorted(fused,key=lambda key:(-fused[key],key))[:budget]
|
|
selected=set(ranked)
|
|
# Whole named source/work scopes are finite, explicit evidence expansion;
|
|
# this is not an assumption that a global top-k result is exhaustive.
|
|
if people:
|
|
selected.update(r['id'] for r in candidates if r['person'] in people)
|
|
rows=[r for r in candidates if r['id'] in selected]
|
|
rows.sort(key=lambda r:(-fused.get(r['id'],0),r['source_path'],r['line']))
|
|
return rows,semantic,{'strategy':'bm25_e5_rrf_assertions','candidate_count':len(rows),
|
|
'lexical_candidates':len(lexical[:budget]),'vector_candidates':len(vector_rank[:budget]),
|
|
'named_scope_expansion':bool(people),'bounded':not people and len(candidates)>len(rows)}
|