Spaces:
Runtime error
Runtime error
Download retrieval_zerogpu.py from rruffi/origin-api: direct link, hf CLI and curl.
- Browser
- Download file 6.48 kB
-
https://huggingface.co/spaces/rruffi/origin-api/resolve/main/retrieval_zerogpu.py
- Command line
-
hf download hf://spaces/rruffi/origin-api/retrieval_zerogpu.py
-
curl -L -o retrieval_zerogpu.py https://huggingface.co/spaces/rruffi/origin-api/resolve/main/retrieval_zerogpu.py
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':[]} | |