origin-api / retrieval_zerogpu.py
rruffi's picture vmcii's picture
Initial Origin Gradio ZeroGPU deployment (#1)
57b1b5d
Raw History Blame Contribute Delete
6.48 kB
"""Origin ZeroGPU retrieval adapter: same scoring/chunking as frozen src/retrieval.py.
Corpus vectors MUST be exported by export_vectors.py from the frozen factory.
Never calculate corpus vectors at Space startup or inside a request.
"""
from __future__ import annotations
import hashlib
import re
import unicodedata
from pathlib import Path
import numpy as np
import pandas as pd
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS
from sklearn.metrics.pairwise import cosine_similarity
MODEL_ID = 'sentence-transformers/paraphrase-multilingual-mpnet-base-v2'
def tokenize(text):
text = unicodedata.normalize('NFKC', str(text)).casefold()
tokens = re.findall(r'(?u)\b\w+\b', text)
return [token for token in tokens if len(token) >= 2 and token not in ENGLISH_STOP_WORDS]
def safe_chunk_text(text, model, overlap=20):
tokenizer = model.tokenizer
token_ids = tokenizer.encode(text, add_special_tokens=False, truncation=False)
special_tokens = tokenizer.num_special_tokens_to_add(pair=False)
chunks = []
start = 0
while start < len(token_ids):
end = min(start + model.max_seq_length - special_tokens, len(token_ids))
while end > start:
chunk = tokenizer.decode(token_ids[start:end], skip_special_tokens=True)
length = len(tokenizer.encode(chunk, add_special_tokens=True, truncation=False))
if length <= model.max_seq_length:
break
end -= 1
if end == start:
raise ValueError('Unable to create a valid chunk')
chunks.append(chunk)
if end == len(token_ids):
break
start = max(start + 1, end - overlap)
return chunks
class OriginRetrieval:
def __init__(self, csv_path: Path, vectors_path: Path):
csv_path, vectors_path = Path(csv_path), Path(vectors_path)
self.df = pd.read_csv(csv_path, encoding='utf-8-sig', keep_default_na=False)
self.df.columns = self.df.columns.str.strip().str.lower()
self.df['department'] = self.df['department'].astype(str).str.strip()
if self.df.shape != (103, 7) or not self.df['project_id'].is_unique:
raise ValueError('Corpus has diverged from frozen 103x7 dataset')
with np.load(vectors_path, allow_pickle=False) as bundle:
ids = bundle['project_ids'].astype(str).tolist()
vectors = bundle['embeddings']
csv_sha = str(bundle['csv_sha256'].item())
model_name = str(bundle['model_name'].item())
if csv_sha != hashlib.sha256(csv_path.read_bytes()).hexdigest():
raise ValueError('Embeddings CSV fingerprint mismatch: re-export vectors')
if model_name != MODEL_ID or ids != self.df['project_id'].astype(str).tolist():
raise ValueError('Embedding model or ID ordering mismatch')
if vectors.shape != (103, 768) or not np.isfinite(vectors).all():
raise ValueError('Invalid embedding matrix')
self.embeddings = vectors.astype(np.float32)
self.project_texts = (
self.df['title'].astype(str).str.strip() + '. ' +
self.df['abstract'].astype(str).str.strip()
).str.replace(r'\s+', ' ', regex=True).str.strip().tolist()
self.bm25 = BM25Okapi([tokenize(t) for t in self.project_texts])
self.records = {str(r['project_id']): r for r in self.df.to_dict('records')}
def query_vector(self, model, title, description):
query = f'{title.strip()}. {description.strip()}'
chunks = safe_chunk_text(query, model)
vectors = model.encode(chunks, batch_size=16, convert_to_numpy=True, normalize_embeddings=True)
vector = vectors.mean(axis=0, keepdims=True)
vector /= np.linalg.norm(vector, axis=1, keepdims=True)
return vector
def search_from_vector(self, query_vector, title, description, lang, top_k=5, rrf_k=60):
if lang not in {'ar', 'en'}: raise ValueError('Unsupported language')
scores = cosine_similarity(query_vector, self.embeddings)[0]
df = self.df
semantic_order = np.argsort(-scores)
if lang == 'ar':
top_indices = semantic_order[:top_k]
return [{'rank':i+1, 'project_id':str(df.iloc[k]['project_id']), 'retrieval_mode':'semantic_ar'}
for i,k in enumerate(top_indices)]
# Equivalent to frozen bm25_search then hybrid_search using full-corpus semantic ranks.
query_tokens = tokenize(f'{title.strip()}. {description.strip()}')
bm25_scores = self.bm25.get_scores(query_tokens)
matched_indices = np.flatnonzero(bm25_scores > 0)
bm25_order = matched_indices[np.argsort(-bm25_scores[matched_indices], kind='stable')]
semantic_ranks = {str(df.iloc[idx]['project_id']): i+1 for i,idx in enumerate(semantic_order)}
bm25_ranks = {str(df.iloc[idx]['project_id']): i+1 for i,idx in enumerate(bm25_order)}
ranking = df[['project_id']].copy()
ranking['semantic_rank'] = ranking['project_id'].map(semantic_ranks)
ranking['bm25_rank'] = ranking['project_id'].map(bm25_ranks)
ranking['rrf_score'] = (1.0 / (rrf_k + ranking['semantic_rank']) +
(1.0 / (rrf_k + ranking['bm25_rank'])).fillna(0.0))
ranking = ranking.sort_values(['rrf_score', 'project_id'], ascending=[False, True]).head(top_k)
mode = 'hybrid_en'
return [{'rank':i+1, 'project_id':str(pid), 'retrieval_mode':mode}
for i,pid in enumerate(ranking['project_id'])]
def public_retrieval(self, ranked, title, description, lang):
results = []
for item in ranked:
row = self.records[item['project_id']]
def both(s): return {'en':str(s), 'ar':str(s)}
results.append({'rank':item['rank'],'project':{
'id':str(row['project_id']), 'title':both(row['title']), 'abstract':both(row['abstract']),
'department':both(row['department']), 'year':str(row['year']),
'keywords':str(row['keywords']), 'sourcePage':int(row['source_page']),
'isFictional':False}})
return {'retrieval': {'idea':{'title':title,'description':description,'language':lang},
'detectedLanguage':lang, 'mode':'arabic_semantic' if lang=='ar' else 'english_hybrid_rrf',
'results':results, 'comparisons':[], 'isDemo':False}, 'suggestions':[]}