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)