Spaces:

Ilaa-1505
/

Explainable-RAG-System

Paused

Update app.py

08679cd verified 6 months ago

7.71 kB

import time
import numpy as np
from flask import Flask, render_template, request, jsonify
from transformers import AutoTokenizer
from sklearn.decomposition import PCA

from src.retrieval.query import (
    retrieve, embed_query, bm25_index, docs_all, metas_all,
    mmr_from_embs, _session
)
from src.generation.generate import generate_answer, build_prompt, build_context

import os
import logging
import warnings

import numpy as np
from src.retrieval.query import _compute_umap

dummy = [np.random.rand(384) for _ in range(6)]
dummy = [e / np.linalg.norm(e) for e in dummy]
_compute_umap(dummy)
print("UMAP warmed up")


# ---------------- ENV + LOGGING ----------------
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

warnings.filterwarnings("ignore")

logging.getLogger("transformers").setLevel(logging.ERROR)
logging.getLogger("sentence_transformers").setLevel(logging.ERROR)
logging.getLogger("urllib3").setLevel(logging.ERROR)

# ---------------- APP ----------------
app = Flask(__name__)

# ---------------- TOKENIZER ----------------
_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-small-en-v1.5")


# ---------------- IDF ----------------
def _build_idf(bm25):
    return {term: max(0.0, float(val)) for term, val in bm25.idf.items()}


_idf_map = _build_idf(bm25_index)
_max_idf = max(_idf_map.values()) if _idf_map else 1.0


# ---------------- PCA ----------------
def _fit_pca(n_components=128):
    import random
    from sentence_transformers import SentenceTransformer

    sample = random.sample(docs_all, min(200, len(docs_all)))
    model = SentenceTransformer("BAAI/bge-small-en-v1.5")
    embs = model.encode(sample, normalize_embeddings=True)

    n_comp = min(n_components, embs.shape[0], embs.shape[1])
    pca = PCA(n_components=n_comp)
    pca.fit(embs)

    return pca


_pca = _fit_pca(128)


# ---------------- ROUTES ----------------

@app.route("/")
def home():
    return render_template("index.html")


# ---------------- QUERY ANALYSIS ----------------
@app.route("/analyze_query", methods=["POST"])
def analyze_query():
    data = request.json
    query = data.get("query", "")

    enc = _tokenizer(query, return_offsets_mapping=True)
    input_ids = enc["input_ids"]
    tokens_raw = _tokenizer.convert_ids_to_tokens(input_ids)

    special = set(_tokenizer.all_special_tokens)
    tokens = [
        {"token": t, "id": int(i)}
        for t, i in zip(tokens_raw, input_ids)
        if t not in special
    ]

    for tok in tokens:
        word = tok["token"].lstrip("##").lower()
        raw_idf = _idf_map.get(word, 0.0)

        tok["idf"] = round(raw_idf, 4)
        tok["idf_normalized"] = round(raw_idf / _max_idf, 4) if _max_idf else 0.0

    idf_vals = [t["idf"] for t in tokens]
    avg_idf = round(sum(idf_vals) / len(idf_vals), 4) if idf_vals else 0.0
    unique_toks = len({t["token"] for t in tokens})

    complexity = round(
        min(1.0, (len(tokens) / 20) * 0.4 + (avg_idf / _max_idf) * 0.6),
        3
    )

    q_emb = embed_query(query)
    projected = _pca.transform(q_emb.reshape(1, -1))[0]

    p_min, p_max = projected.min(), projected.max()

    if p_max != p_min:
        normed = ((projected - p_min) / (p_max - p_min) * 2 - 1).tolist()
    else:
        normed = [0.0] * len(projected)

    return jsonify({
        "tokens": tokens,
        "embedding": [round(v, 4) for v in normed],
        "stats": {
            "token_count": len(tokens),
            "unique_tokens": unique_toks,
            "avg_idf": avg_idf,
            "complexity": complexity,
        }
    })


# ---------------- MMR RERUN ----------------
@app.route("/mmr_rerun", methods=["POST"])
def mmr_rerun():
    if _session["query_emb"] is None:
        return jsonify({"error": "No active session. Run a query first."}), 400

    data = request.json
    lambda_param = float(data.get("lambda", 0.7))
    lambda_param = max(0.0, min(1.0, lambda_param))

    selected = mmr_from_embs(
        _session["query_emb"],
        _session["doc_indices"],
        _session["embs"],
        k=10,
        lambda_param=lambda_param
    )

    return jsonify({
        "selected_indices": selected,
        "selected_local": [
            _session["doc_indices"].index(s) for s in selected
        ],
        "lambda": lambda_param,
    })


# ---------------- MAIN RAG ----------------
@app.route("/ask", methods=["POST"])
def ask():
    data = request.json
    query = data.get("query")

    print(f"\n[API QUERY]: {query}\n")

    # -------- RETRIEVE --------
    t0 = time.perf_counter()
    results, debug = retrieve(query)
    t_retrieve = time.perf_counter() - t0

    # -------- SORT --------
    results = sorted(results, key=lambda x: (
        x["meta"].get("chunk_id", 0),
        x["meta"].get("global_chunk_id", 0)
    ))

    docs = [r["text"] for r in results]
    metas = [r["meta"] for r in results]
    raw_scores = [float(r["rerank_score"]) for r in results]

    # -------- CONTEXT --------
    context = build_context(docs, metas, raw_scores)

    # -------- LLM --------
    t1 = time.perf_counter()
    prompt = build_prompt(query, context)
    answer = generate_answer(prompt)
    t_llm = time.perf_counter() - t1

    # -------- SOURCES --------
    sources = [
        {
            "title": meta.get("title", "Source"),
            "url": meta.get("url", "")
        }
        for meta in metas
    ]

    # -------- TIMINGS --------
    stage_timings = debug.get("timings", {})
    stage_timings["llm"] = round(t_llm * 1000)
    stage_timings["total"] = round((t_retrieve + t_llm) * 1000)

    # -------- COMPARISON --------
    score_lookup = debug.get("score_lookup", {})
    full_rerank = debug.get("rerank_full", [])

    hybrid_order = {
        int(k): rank for rank, k in enumerate(score_lookup.keys())
    }

    comparison_rows = []

    for post_rank, (idx, rerank_score) in enumerate(full_rerank):
        idx = int(idx)
        sk = score_lookup.get(str(idx), [0, 0, 0])

        pre_rank = hybrid_order.get(idx, post_rank)

        comparison_rows.append({
            "idx": idx,
            "pre_rank": pre_rank,
            "post_rank": post_rank,
            "rank_delta": pre_rank - post_rank,
            "vector_score": round(float(sk[0]), 4),
            "bm25_score": round(float(sk[1]), 4),
            "hybrid_score": round(float(sk[2]), 4),
            "rerank_score": round(float(rerank_score), 4),
            "passed_threshold": float(rerank_score) >= 0.3,
            "text_preview": " ".join(
                docs_all[idx]
                .replace("passage: ", "")
                .strip()
                .lstrip("`")
                .split()
            )[:120],
            "text_full": docs_all[idx].replace("passage: ", ""),
            "title": metas_all[idx].get("title", ""),
        })

    # -------- RESPONSE --------
    return jsonify({
        "answer": answer,
        "sources": sources,
        "chunks": docs,
        "scores": raw_scores,
        "raw_scores": raw_scores,
        "debug": debug,
        "timings": stage_timings,
        "comparison_rows": comparison_rows,
        "mmr_data": {
            "umap_coords": debug.get("umap_coords"),
            "sim_matrix": debug.get("sim_matrix"),
            "doc_indices": debug.get("doc_indices", []),
            "sims": debug.get("sims", []),
            "doc_previews": debug.get("doc_previews", []),
            "mmr_selected": debug.get("mmr_selected", []),
            "no_mmr_selected": debug.get("no_mmr_selected", []),
        }
    })


# ---------------- RUN ----------------
if __name__ == "__main__":
    app.run(debug=True)

	import time
	import numpy as np
	from flask import Flask, render_template, request, jsonify
	from transformers import AutoTokenizer
	from sklearn.decomposition import PCA

	from src.retrieval.query import (
	retrieve, embed_query, bm25_index, docs_all, metas_all,
	mmr_from_embs, _session
	)
	from src.generation.generate import generate_answer, build_prompt, build_context

	import os
	import logging
	import warnings

	import numpy as np
	from src.retrieval.query import _compute_umap

	dummy = [np.random.rand(384) for _ in range(6)]
	dummy = [e / np.linalg.norm(e) for e in dummy]
	_compute_umap(dummy)
	print("UMAP warmed up")


	# ---------------- ENV + LOGGING ----------------
	os.environ["TRANSFORMERS_VERBOSITY"] = "error"
	os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
	os.environ["TOKENIZERS_PARALLELISM"] = "false"

	warnings.filterwarnings("ignore")

	logging.getLogger("transformers").setLevel(logging.ERROR)
	logging.getLogger("sentence_transformers").setLevel(logging.ERROR)
	logging.getLogger("urllib3").setLevel(logging.ERROR)

	# ---------------- APP ----------------
	app = Flask(__name__)

	# ---------------- TOKENIZER ----------------
	_tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-small-en-v1.5")


	# ---------------- IDF ----------------
	def _build_idf(bm25):
	return {term: max(0.0, float(val)) for term, val in bm25.idf.items()}


	_idf_map = _build_idf(bm25_index)
	_max_idf = max(_idf_map.values()) if _idf_map else 1.0


	# ---------------- PCA ----------------
	def _fit_pca(n_components=128):
	import random
	from sentence_transformers import SentenceTransformer

	sample = random.sample(docs_all, min(200, len(docs_all)))
	model = SentenceTransformer("BAAI/bge-small-en-v1.5")
	embs = model.encode(sample, normalize_embeddings=True)

	n_comp = min(n_components, embs.shape[0], embs.shape[1])
	pca = PCA(n_components=n_comp)
	pca.fit(embs)

	return pca


	_pca = _fit_pca(128)


	# ---------------- ROUTES ----------------

	@app.route("/")
	def home():
	return render_template("index.html")


	# ---------------- QUERY ANALYSIS ----------------
	@app.route("/analyze_query", methods=["POST"])
	def analyze_query():
	data = request.json
	query = data.get("query", "")

	enc = _tokenizer(query, return_offsets_mapping=True)
	input_ids = enc["input_ids"]
	tokens_raw = _tokenizer.convert_ids_to_tokens(input_ids)

	special = set(_tokenizer.all_special_tokens)
	tokens = [
	{"token": t, "id": int(i)}
	for t, i in zip(tokens_raw, input_ids)
	if t not in special
	]

	for tok in tokens:
	word = tok["token"].lstrip("##").lower()
	raw_idf = _idf_map.get(word, 0.0)

	tok["idf"] = round(raw_idf, 4)
	tok["idf_normalized"] = round(raw_idf / _max_idf, 4) if _max_idf else 0.0

	idf_vals = [t["idf"] for t in tokens]
	avg_idf = round(sum(idf_vals) / len(idf_vals), 4) if idf_vals else 0.0
	unique_toks = len({t["token"] for t in tokens})

	complexity = round(
	min(1.0, (len(tokens) / 20) * 0.4 + (avg_idf / _max_idf) * 0.6),
	3
	)

	q_emb = embed_query(query)
	projected = _pca.transform(q_emb.reshape(1, -1))[0]

	p_min, p_max = projected.min(), projected.max()

	if p_max != p_min:
	normed = ((projected - p_min) / (p_max - p_min) * 2 - 1).tolist()
	else:
	normed = [0.0] * len(projected)

	return jsonify({
	"tokens": tokens,
	"embedding": [round(v, 4) for v in normed],
	"stats": {
	"token_count": len(tokens),
	"unique_tokens": unique_toks,
	"avg_idf": avg_idf,
	"complexity": complexity,
	}
	})


	# ---------------- MMR RERUN ----------------
	@app.route("/mmr_rerun", methods=["POST"])
	def mmr_rerun():
	if _session["query_emb"] is None:
	return jsonify({"error": "No active session. Run a query first."}), 400

	data = request.json
	lambda_param = float(data.get("lambda", 0.7))
	lambda_param = max(0.0, min(1.0, lambda_param))

	selected = mmr_from_embs(
	_session["query_emb"],
	_session["doc_indices"],
	_session["embs"],
	k=10,
	lambda_param=lambda_param
	)

	return jsonify({
	"selected_indices": selected,
	"selected_local": [
	_session["doc_indices"].index(s) for s in selected
	],
	"lambda": lambda_param,
	})


	# ---------------- MAIN RAG ----------------
	@app.route("/ask", methods=["POST"])
	def ask():
	data = request.json
	query = data.get("query")

	print(f"\n[API QUERY]: {query}\n")

	# -------- RETRIEVE --------
	t0 = time.perf_counter()
	results, debug = retrieve(query)
	t_retrieve = time.perf_counter() - t0

	# -------- SORT --------
	results = sorted(results, key=lambda x: (
	x["meta"].get("chunk_id", 0),
	x["meta"].get("global_chunk_id", 0)
	))

	docs = [r["text"] for r in results]
	metas = [r["meta"] for r in results]
	raw_scores = [float(r["rerank_score"]) for r in results]

	# -------- CONTEXT --------
	context = build_context(docs, metas, raw_scores)

	# -------- LLM --------
	t1 = time.perf_counter()
	prompt = build_prompt(query, context)
	answer = generate_answer(prompt)
	t_llm = time.perf_counter() - t1

	# -------- SOURCES --------
	sources = [
	{
	"title": meta.get("title", "Source"),
	"url": meta.get("url", "")
	}
	for meta in metas
	]

	# -------- TIMINGS --------
	stage_timings = debug.get("timings", {})
	stage_timings["llm"] = round(t_llm * 1000)
	stage_timings["total"] = round((t_retrieve + t_llm) * 1000)

	# -------- COMPARISON --------
	score_lookup = debug.get("score_lookup", {})
	full_rerank = debug.get("rerank_full", [])

	hybrid_order = {
	int(k): rank for rank, k in enumerate(score_lookup.keys())
	}

	comparison_rows = []

	for post_rank, (idx, rerank_score) in enumerate(full_rerank):
	idx = int(idx)
	sk = score_lookup.get(str(idx), [0, 0, 0])

	pre_rank = hybrid_order.get(idx, post_rank)

	comparison_rows.append({
	"idx": idx,
	"pre_rank": pre_rank,
	"post_rank": post_rank,
	"rank_delta": pre_rank - post_rank,
	"vector_score": round(float(sk[0]), 4),
	"bm25_score": round(float(sk[1]), 4),
	"hybrid_score": round(float(sk[2]), 4),
	"rerank_score": round(float(rerank_score), 4),
	"passed_threshold": float(rerank_score) >= 0.3,
	"text_preview": " ".join(
	docs_all[idx]
	.replace("passage: ", "")
	.strip()
	.lstrip("`")
	.split()
	)[:120],
	"text_full": docs_all[idx].replace("passage: ", ""),
	"title": metas_all[idx].get("title", ""),
	})

	# -------- RESPONSE --------
	return jsonify({
	"answer": answer,
	"sources": sources,
	"chunks": docs,
	"scores": raw_scores,
	"raw_scores": raw_scores,
	"debug": debug,
	"timings": stage_timings,
	"comparison_rows": comparison_rows,
	"mmr_data": {
	"umap_coords": debug.get("umap_coords"),
	"sim_matrix": debug.get("sim_matrix"),
	"doc_indices": debug.get("doc_indices", []),
	"sims": debug.get("sims", []),
	"doc_previews": debug.get("doc_previews", []),
	"mmr_selected": debug.get("mmr_selected", []),
	"no_mmr_selected": debug.get("no_mmr_selected", []),
	}
	})


	# ---------------- RUN ----------------
	if __name__ == "__main__":
	app.run(debug=True)