Spaces:

PranilB
/

Question_Answering_Classification

Sleeping

Question_Answering_Classification / rag /retriever.py

Pranil Bankar

Deploy Application

4d67ec7 5 months ago

3.25 kB

import os
from dotenv import load_dotenv
from supabase import create_client, Client
from sentence_transformers import SentenceTransformer

# Load .env for local development (Railway injects env vars directly)
_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
load_dotenv(dotenv_path=os.path.join(_root, '.env'), override=False)
load_dotenv(dotenv_path=os.path.join(_root, 'pdf_to_embedding', '.env'), override=False)

# ==============================
# CONFIG
# ==============================
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5")
TOP_K = 5                    # Number of chunks to retrieve
SIMILARITY_THRESHOLD = 0.1   # Low floor — let chunks through, LLM + verifier will judge quality

# ==============================
# SINGLETON CLIENTS
# ==============================
_supabase_client: Client = None
_embedding_model: SentenceTransformer = None


def get_supabase() -> Client:
    global _supabase_client
    if _supabase_client is None:
        if not SUPABASE_URL or not SUPABASE_KEY:
            raise ValueError("Missing SUPABASE_URL or SUPABASE_KEY in .env file.")
        _supabase_client = create_client(SUPABASE_URL, SUPABASE_KEY)
    return _supabase_client


def get_embedding_model() -> SentenceTransformer:
    global _embedding_model
    if _embedding_model is None:
        print(f"[Retriever] Loading embedding model: {EMBEDDING_MODEL_NAME}")
        _embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
    return _embedding_model


# ==============================
# RETRIEVAL FUNCTION
# ==============================
def retrieve_chunks(
    question: str,
    chapter_filter: str = None,
    top_k: int = TOP_K,
) -> list[dict]:
    """
    Embeds the question and retrieves the top-k most similar NCERT chunks
    from Supabase pgvector, optionally filtered by chapter.
    Returns a list of dicts: {chunk_id, chapter, section_title, text_content, similarity}
    """
    model = get_embedding_model()
    supabase = get_supabase()

    # Embed the question
    query_vector = model.encode(question, normalize_embeddings=True).tolist()

    # Call the Postgres similarity search function
    try:
        response = supabase.rpc(
            "match_ncert_chunks",
            {
                "query_embedding": query_vector,
                "match_threshold": SIMILARITY_THRESHOLD,
                "match_count": top_k,
                "filter_chapter": chapter_filter,
            },
        ).execute()
    except Exception as e:
        print(f"[Retriever] Error calling match_ncert_chunks: {e}")
        return []

    results = response.data or []

    if not results:
        print("[Retriever] Warning: No chunks found above the similarity threshold.")

    return results


if __name__ == "__main__":
    # Quick smoke test
    question = "What is the role of mitochondria in a cell?"
    chunks = retrieve_chunks(question, top_k=3)
    print(f"\nRetrieved {len(chunks)} chunks:")
    for i, c in enumerate(chunks, 1):
        print(f"\n[{i}] ({c['similarity']:.3f}) {c['chapter']} > {c['section_title']}")
        print(f"    {c['text_content'][:120]}...")

	import os
	from dotenv import load_dotenv
	from supabase import create_client, Client
	from sentence_transformers import SentenceTransformer

	# Load .env for local development (Railway injects env vars directly)
	_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
	load_dotenv(dotenv_path=os.path.join(_root, '.env'), override=False)
	load_dotenv(dotenv_path=os.path.join(_root, 'pdf_to_embedding', '.env'), override=False)

	# ==============================
	# CONFIG
	# ==============================
	SUPABASE_URL = os.getenv("SUPABASE_URL")
	SUPABASE_KEY = os.getenv("SUPABASE_KEY")
	EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5")
	TOP_K = 5 # Number of chunks to retrieve
	SIMILARITY_THRESHOLD = 0.1 # Low floor — let chunks through, LLM + verifier will judge quality

	# ==============================
	# SINGLETON CLIENTS
	# ==============================
	_supabase_client: Client = None
	_embedding_model: SentenceTransformer = None


	def get_supabase() -> Client:
	global _supabase_client
	if _supabase_client is None:
	if not SUPABASE_URL or not SUPABASE_KEY:
	raise ValueError("Missing SUPABASE_URL or SUPABASE_KEY in .env file.")
	_supabase_client = create_client(SUPABASE_URL, SUPABASE_KEY)
	return _supabase_client


	def get_embedding_model() -> SentenceTransformer:
	global _embedding_model
	if _embedding_model is None:
	print(f"[Retriever] Loading embedding model: {EMBEDDING_MODEL_NAME}")
	_embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
	return _embedding_model


	# ==============================
	# RETRIEVAL FUNCTION
	# ==============================
	def retrieve_chunks(
	question: str,
	chapter_filter: str = None,
	top_k: int = TOP_K,
	) -> list[dict]:
	"""
	Embeds the question and retrieves the top-k most similar NCERT chunks
	from Supabase pgvector, optionally filtered by chapter.

	Returns a list of dicts: {chunk_id, chapter, section_title, text_content, similarity}
	"""
	model = get_embedding_model()
	supabase = get_supabase()

	# Embed the question
	query_vector = model.encode(question, normalize_embeddings=True).tolist()

	# Call the Postgres similarity search function
	try:
	response = supabase.rpc(
	"match_ncert_chunks",
	{
	"query_embedding": query_vector,
	"match_threshold": SIMILARITY_THRESHOLD,
	"match_count": top_k,
	"filter_chapter": chapter_filter,
	},
	).execute()
	except Exception as e:
	print(f"[Retriever] Error calling match_ncert_chunks: {e}")
	return []

	results = response.data or []

	if not results:
	print("[Retriever] Warning: No chunks found above the similarity threshold.")

	return results


	if __name__ == "__main__":
	# Quick smoke test
	question = "What is the role of mitochondria in a cell?"
	chunks = retrieve_chunks(question, top_k=3)
	print(f"\nRetrieved {len(chunks)} chunks:")
	for i, c in enumerate(chunks, 1):
	print(f"\n[{i}] ({c['similarity']:.3f}) {c['chapter']} > {c['section_title']}")
	print(f" {c['text_content'][:120]}...")