r"""
J.R. Kantor Research System — Hugging Face Spaces App (v2)
==========================================================
Powered by:
- Embeddings: nomic-ai/nomic-embed-text-v1.5 (768-dim, 8192 token context)
- LLM: Groq llama-3.3-70b-versatile
- Vector DB: FAISS (local, cosine similarity via inner product)
- Frontend: Gradio with dark red academic theme
- Index: Page-aware chunking with page number metadata
"""
import gradio as gr
import os
import re
import json
import pickle
import numpy as np
import faiss
from groq import Groq
from sentence_transformers import SentenceTransformer, CrossEncoder
# ── Config ──────────────────────────────────────────────────────
DB_DIR = "./database"
EMBEDDING_MODEL = "nomic-ai/nomic-embed-text-v1.5"
RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
# ── Initialize ──────────────────────────────────────────────────
print("Loading embedding model...")
model = SentenceTransformer(EMBEDDING_MODEL, trust_remote_code=True)
print(f"Model loaded: {EMBEDDING_MODEL} (dim={model.get_sentence_embedding_dimension()})")
print("Loading reranker model...")
reranker = CrossEncoder(RERANKER_MODEL, max_length=512)
print(f"Reranker loaded: {RERANKER_MODEL}")
print("Loading FAISS index...")
index = faiss.read_index(os.path.join(DB_DIR, "index.faiss"))
print(f"FAISS index: {index.ntotal} vectors")
print("Loading metadata...")
with open(os.path.join(DB_DIR, "index_data.pkl"), "rb") as f:
data = pickle.load(f)
texts = data["texts"]
metadatas = data["metadatas"]
print(f"Metadata loaded: {len(texts)} chunks")
# Load summary for stats
summary_path = os.path.join(DB_DIR, "index_summary.json")
if os.path.exists(summary_path):
with open(summary_path) as f:
index_summary = json.load(f)
else:
index_summary = {"total_chunks": len(texts), "total_entries": "?", "types": {}}
print("Initializing Groq client...")
groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
print("System ready!")
# ── Search & Answer ─────────────────────────────────────────────
TYPE_LABELS = {
"Book": "Book",
"Article": "Article",
"Review": "Book Review",
"Observer": "Observer Column",
}
def format_page_ref(meta):
"""Format page reference from metadata."""
start = meta.get("start_page")
end = meta.get("end_page")
if not start:
return ""
if start == end or not end:
return f"p. {start}"
return f"pp. {start}–{end}"
def format_citation(meta):
"""Build clean APA-style citation with page numbers."""
title = meta.get("title", "Unknown").strip()
year = meta.get("year", "n.d.")
doc_type = meta.get("type", "")
label = TYPE_LABELS.get(doc_type, doc_type)
page_ref = format_page_ref(meta)
citation = f"Kantor, J.R. ({year}). *{title}* [{label}]"
if page_ref:
citation += f", {page_ref}"
return citation
def format_citation_for_llm(meta):
"""Shorter citation for the LLM context window."""
title = meta.get("title", "Unknown").strip()
year = meta.get("year", "n.d.")
page_ref = format_page_ref(meta)
cite = f"Kantor ({year}), \"{title}\""
if page_ref:
cite += f", {page_ref}"
return cite
def is_valid_chunk(text):
"""Filter out chunks that are just punctuation, numbers, or metadata."""
if not text or len(text.strip()) < 30:
return False
alphanumeric = re.sub(r"[^a-zA-Z]", "", text)
return len(alphanumeric) >= 20
def is_metadata_heavy(text, meta):
"""Detect chunks that are mostly title/abstract/journal metadata.
These match keywords well but contain little substantive content."""
t = text.lower()
if 'requests for reprints' in t:
return True
# Page 1 chunks with telltale signs of abstracts/headers
if meta.get('start_page', 0) == 1:
signals = 0
if 'abstract' in t: signals += 2
if 'resumen' in t: signals += 2
if 'vol.' in t or 'vol ' in t: signals += 1
if 'pp.' in t or 'pp ' in t: signals += 1
if 'university' in t and 'indiana' in t.lower(): signals += 1
if 'revista' in t or 'journal' in t: signals += 1
if signals >= 3:
return True
return False
def is_reference_list(text):
"""Detect chunks that are bibliography / reference-list pages.
They are dense with citations and match keywords, but contain no
substantive discussion."""
# Repeated author citations ("Kantor, J. R." / "Kantor, J.R.")
author_cites = len(re.findall(r"Kantor,\s*J\.?\s*R\.?", text))
if author_cites >= 3:
return True
# Journal citation pattern: "1922, 19, 38-49" (year, volume, pages)
journal_cites = len(re.findall(r"\b1[89]\d{2},\s*\d+,\s*\d+", text))
if journal_cites >= 3:
return True
# High density of years (typical of reference pages)
years = len(re.findall(r"\b1[89]\d{2}\b", text))
if years >= 6 and years / max(len(text) / 100.0, 1.0) > 1.0:
return True
return False
def word_quality(text):
"""Fraction of tokens that look like real words. Low values indicate
badly OCR'd pages that pollute results."""
tokens = re.findall(r"[A-Za-z]+", text)
if not tokens:
return 0.0
good = sum(1 for t in tokens
if len(t) >= 2 and (t.islower() or t.istitle()
or (t.isupper() and len(t) <= 4)))
return good / len(tokens)
def search_faiss(query, k=40):
"""Search FAISS with nomic embeddings. Skips junk chunks
(reference lists, title pages, metadata-heavy or badly OCR'd pages)."""
prefixed = f"search_query: {query}"
query_vec = model.encode([prefixed], normalize_embeddings=True).astype("float32")
scores, indices = index.search(query_vec, k)
results = []
for score, idx in zip(scores[0], indices[0]):
if idx < 0 or idx >= len(texts):
continue
text = texts[idx]
meta = metadatas[idx]
if not is_valid_chunk(text):
continue
if is_reference_list(text):
continue
if is_metadata_heavy(text, meta):
continue
if word_quality(text) < 0.78:
continue
results.append({
"text": text,
"meta": meta,
"score": float(score),
})
return results
def rerank_results(query, results, keep=16):
"""Re-score candidates with a cross-encoder and blend with the
embedding score. The blend keeps the strengths of both: embeddings
capture topic, the cross-encoder captures actual relevance."""
if len(results) < 2:
return results
pairs = [(query, r["text"]) for r in results]
ce_scores = np.array(reranker.predict(pairs), dtype=float)
emb_scores = np.array([r["score"] for r in results], dtype=float)
def z(x):
return (x - x.mean()) / (x.std() + 1e-9)
final = 0.6 * z(emb_scores) + 0.4 * z(ce_scores)
for r, s in zip(results, final):
r["rerank_score"] = float(s)
results.sort(key=lambda r: -r["rerank_score"])
return results[:keep]
def deduplicate_results(results, max_results=8):
"""Deduplicate by content similarity AND by title.
- Skip chunks with near-identical text (catches journal issue duplicates)
- Allow up to 2 chunks from the same work"""
seen_text_starts = set()
title_counts = {}
unique = []
for r in results:
# Content dedup: skip if first 150 chars match something already selected
text_key = r["text"][:150].strip().lower()
if text_key in seen_text_starts:
continue
seen_text_starts.add(text_key)
# Title dedup: max 2 per work
title = r["meta"].get("title", "")
count = title_counts.get(title, 0)
if count >= 2:
continue
title_counts[title] = count + 1
unique.append(r)
if len(unique) >= max_results:
break
return unique
def search_and_answer(query):
"""Main RAG pipeline: search -> context -> LLM -> answer."""
if not query or len(query.strip()) < 3:
return "Please enter a research question (at least 3 characters).", ""
try:
results = search_faiss(query, k=40)
reranked = rerank_results(query, results, keep=16)
unique_results = deduplicate_results(reranked, max_results=8)
if not unique_results:
return "No relevant passages found for this query. Try rephrasing your question.", ""
# Build context for LLM
context = ""
source_refs = ""
sources_display = []
for i, r in enumerate(unique_results, 1):
llm_citation = format_citation_for_llm(r["meta"])
display_citation = format_citation(r["meta"])
context += f"\n[Source {i}: {llm_citation}]\n{r['text']}\n"
source_refs += f"- Source {i}: {llm_citation}\n"
sources_display.append({
"num": i,
"citation": display_citation,
"text": r["text"],
})
# LLM call
response = groq_client.chat.completions.create(
model="llama-3.3-70b-versatile",
messages=[
{
"role": "system",
"content": f"""You are a research assistant specializing in J.R. Kantor's interbehavioral psychology. Your audience includes students, researchers, and academics exploring Kantor's work.
Answer the question based ONLY on the provided source excerpts from Kantor's works.
Guidelines:
1. Cite every substantive claim using [Source X] notation.
2. Synthesize across multiple sources when they address the same topic.
3. Use clear academic prose — be thorough but accessible.
4. If different sources show how Kantor's thinking evolved over time, note this.
5. Do NOT cite authors other than Kantor unless they appear in the excerpts.
6. Do NOT invent or fabricate any references.
7. If the excerpts don't fully address the question, state what IS covered and note the limitation.
Available sources:
{source_refs}"""
},
{
"role": "user",
"content": f"Excerpts from Kantor's works:\n{context}\n\nQuestion: {query}"
}
],
temperature=0.4,
max_tokens=1200,
frequency_penalty=0.3,
)
answer = response.choices[0].message.content
# Build sources display
sources_md = ""
for s in sources_display:
sources_md += f"**Source {s['num']}:** {s['citation']}\n\n"
display_text = s["text"][:600] + "..." if len(s["text"]) > 600 else s["text"]
sources_md += f"View passage
\n\n{display_text}\n\n