Spaces:
Running on Zero
Running on Zero
File size: 4,821 Bytes
6c3069e a3037d6 acbca58 a3037d6 6c3069e a3037d6 6c3069e a3037d6 6c3069e a3037d6 6c3069e a3037d6 6c3069e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Embedding and FAISS retrieval — REVISED v3."""
import os
import pickle
from typing import List, Dict
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from config import (
EMBEDDING_MODEL, EMBEDDING_DIM, INDEX_DIR, TOP_K, TOP_K_FETCH, get_boosted_sections
)
class PortfolioRetriever:
"""Dual-tier retriever: FAISS semantic + keyword fallback."""
def __init__(self, build: bool = False):
self.encoder = SentenceTransformer(EMBEDDING_MODEL, device="cpu")
self.index = None
self.chunks = []
if build or not self._index_exists():
self._build_index()
else:
self._load_index()
def _index_exists(self) -> bool:
return os.path.exists(os.path.join(INDEX_DIR, "faiss.index"))
def _build_index(self):
"""Build FAISS index from KB chunks."""
from ingestion import parse_kb
print("Building FAISS index...")
os.makedirs(INDEX_DIR, exist_ok=True)
self.chunks = parse_kb()
texts = [c["text"] for c in self.chunks]
print(f"Encoding {len(texts)} chunks with {EMBEDDING_MODEL}...")
embeddings = self.encoder.encode(
texts,
show_progress_bar=True,
normalize_embeddings=True
)
self.index = faiss.IndexFlatIP(EMBEDDING_DIM)
self.index.add(np.array(embeddings).astype('float32'))
faiss.write_index(self.index, os.path.join(INDEX_DIR, "faiss.index"))
with open(os.path.join(INDEX_DIR, "chunks.pkl"), 'wb') as f:
pickle.dump(self.chunks, f)
print(f"Index saved: {len(self.chunks)} chunks, {EMBEDDING_DIM}d")
def _load_index(self):
"""Load pre-built index."""
print("Loading pre-built FAISS index...")
self.index = faiss.read_index(os.path.join(INDEX_DIR, "faiss.index"))
with open(os.path.join(INDEX_DIR, "chunks.pkl"), 'rb') as f:
self.chunks = pickle.load(f)
print(f"Loaded {len(self.chunks)} chunks")
def retrieve(self, query: str, persona: str = "general", k: int = TOP_K) -> List[Dict]:
"""Retrieve relevant chunks with keyword fallback boosting."""
query_vec = self.encoder.encode([query], normalize_embeddings=True)
query_vec = np.array(query_vec).astype('float32')
scores, indices = self.index.search(query_vec, min(TOP_K_FETCH, len(self.chunks)))
results = []
seen_ids = set()
boosted_sections = get_boosted_sections(query)
for score, idx in zip(scores[0], indices[0]):
if idx < 0 or idx >= len(self.chunks):
continue
chunk = self.chunks[idx].copy()
chunk_id = chunk["id"]
if chunk_id in seen_ids:
continue
seen_ids.add(chunk_id)
chunk["score"] = float(score)
if chunk["section_type"] in boosted_sections:
chunk["score"] *= 1.5
if persona == "recruiter" and chunk["section_type"] in [
"identity", "education", "skills", "skills_summary",
"experience", "experience_summary_v2", "certifications",
"kaggle_summary"
]:
chunk["score"] *= 1.15
elif persona == "technical" and chunk["section_type"] in [
"architecture_deep_dive_axiomis", "projects", "skills", "skills_summary",
"publications", "publications_summary", "philosophy",
"clinical_ai_projects", "marketing_ai_projects", "supply_chain_projects",
"energy_projects", "computational_biology_projects", "finance_regtech_projects",
"education_projects", "neuro_symbolic_projects", "projects_master_list"
]:
chunk["score"] *= 1.15
results.append(chunk)
results.sort(key=lambda x: x["score"], reverse=True)
final_results = []
section_counts = {}
for r in results:
st = r["section_type"]
section_counts[st] = section_counts.get(st, 0) + 1
if section_counts[st] <= 3:
final_results.append(r)
if len(final_results) >= k:
break
return final_results
def get_stats(self) -> Dict:
from collections import Counter
dist = Counter(c["section_type"] for c in self.chunks)
return {
"total_chunks": len(self.chunks),
"embedding_dim": EMBEDDING_DIM,
"encoder": EMBEDDING_MODEL,
"section_distribution": dict(dist)
} |