Spaces:
Paused
Paused
| import os | |
| import re | |
| import string | |
| import warnings | |
| from collections import Counter | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| import faiss | |
| import gradio as gr | |
| import torch | |
| from sentence_transformers import SentenceTransformer | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| warnings.filterwarnings("ignore") | |
| TOOL_NAME = "SourceSure" | |
| TOOL_TAGLINE = "Grounded answers with a clear source trail." | |
| # Config: keep these aligned with the Kaggle notebook defaults. | |
| CHECKPOINT = "selfrag/selfrag_llama2_7b" | |
| DELTA = 0.2 | |
| BEAM_WIDTH = 2 | |
| K_PASSAGES = 5 | |
| MAX_NEW_TOKS = 100 | |
| MAX_SEGMENTS = 5 | |
| W_REL = 1.0 | |
| W_SUP = 1.0 | |
| W_USE = 0.5 | |
| ABSTENTION_ISREL_THRESHOLD = 0.20 | |
| ABSTENTION_MIN_QUERY_COVERAGE = 0.20 | |
| # Optional runtime tweak for Hugging Face Spaces. | |
| # Default stays aligned with the notebook: float16 on GPU. | |
| LOAD_IN_4BIT = os.getenv("LOAD_IN_4BIT", "0") == "1" | |
| class RetrieveToken: | |
| YES = "[Retrieve=Yes]" | |
| NO = "[Retrieve=No]" | |
| CONTINUE = "[Retrieve=Continue]" | |
| ALL = [YES, NO, CONTINUE] | |
| class IsRelToken: | |
| RELEVANT = "[IsRel=Relevant]" | |
| IRRELEVANT = "[IsRel=Irrelevant]" | |
| ALL = [RELEVANT, IRRELEVANT] | |
| class IsSupportToken: | |
| FULLY = "[IsSupport=Fully Supported]" | |
| PARTIALLY = "[IsSupport=Partially Supported]" | |
| NO = "[IsSupport=No Support]" | |
| ALL = [FULLY, PARTIALLY, NO] | |
| class IsUseToken: | |
| FIVE = "[IsUse=5]" | |
| FOUR = "[IsUse=4]" | |
| THREE = "[IsUse=3]" | |
| TWO = "[IsUse=2]" | |
| ONE = "[IsUse=1]" | |
| ALL = [FIVE, FOUR, THREE, TWO, ONE] | |
| WEIGHTS = {1: -1.0, 2: -0.5, 3: 0.0, 4: 0.5, 5: 1.0} | |
| ALL_REFLECTION_TOKENS = ( | |
| RetrieveToken.ALL + IsRelToken.ALL + IsSupportToken.ALL + IsUseToken.ALL | |
| ) | |
| DOCUMENTS = [ | |
| { | |
| "id": "doc_00", | |
| "title": "RAG Overview", | |
| "text": ( | |
| "Retrieval-Augmented Generation (RAG) improves factual correctness in large " | |
| "language models by combining a neural generator with a non-parametric retrieval " | |
| "system. The model is given relevant documents at inference time instead of " | |
| "relying solely on memorised parametric knowledge. RAG was introduced by Lewis " | |
| "et al. in 2020." | |
| ), | |
| }, | |
| { | |
| "id": "doc_01", | |
| "title": "Vanilla RAG Limitations", | |
| "text": ( | |
| "Vanilla RAG always retrieves a fixed number of top-k documents regardless of " | |
| "whether retrieval is actually necessary for the query. This indiscriminate " | |
| "retrieval can introduce noise when the documents are not relevant to the " | |
| "question, and the model may produce outputs that are inconsistent with " | |
| "retrieved passages." | |
| ), | |
| }, | |
| { | |
| "id": "doc_02", | |
| "title": "SELF-RAG", | |
| "text": ( | |
| "SELF-RAG is a framework that trains a single language model to retrieve, " | |
| "generate, and critique text through self-reflection. It uses four special " | |
| "reflection tokens: Retrieve (yes/no/continue), IsRel (relevant/irrelevant), " | |
| "IsSupport (fully/partially/no support), and IsUse (1 to 5). These tokens " | |
| "allow the model to decide adaptively when to retrieve and to evaluate the " | |
| "quality of its own generation at inference time without any external critic " | |
| "model." | |
| ), | |
| }, | |
| { | |
| "id": "doc_03", | |
| "title": "CRAG - Corrective RAG", | |
| "text": ( | |
| "Corrective Retrieval-Augmented Generation (CRAG) addresses the problem of " | |
| "inaccurate or irrelevant retrieved documents by adding a lightweight evaluator " | |
| "that grades each retrieved document's relevance. When all retrieved documents " | |
| "are judged irrelevant, CRAG triggers a corrective web search to obtain better " | |
| "evidence before generation." | |
| ), | |
| }, | |
| { | |
| "id": "doc_04", | |
| "title": "RAGAs Evaluation", | |
| "text": ( | |
| "RAGAs is an automated evaluation framework for retrieval-augmented generation " | |
| "systems. It provides metrics including faithfulness, which measures whether " | |
| "the generated answer is supported by the retrieved context, and answer " | |
| "relevance, which measures whether the answer addresses the original question." | |
| ), | |
| }, | |
| { | |
| "id": "doc_05", | |
| "title": "Hallucination in LLMs", | |
| "text": ( | |
| "Hallucination in large language models refers to the generation of fluent and " | |
| "confident statements that are not supported by evidence or are factually " | |
| "incorrect. RAG reduces hallucination by grounding generation in retrieved " | |
| "documents, but does not eliminate it entirely when the model ignores or " | |
| "misuses the retrieved context." | |
| ), | |
| }, | |
| { | |
| "id": "doc_06", | |
| "title": "Adaptive Retrieval in SELF-RAG", | |
| "text": ( | |
| "SELF-RAG uses an adaptive retrieval threshold delta to control how often " | |
| "retrieval is triggered. The model predicts a Retrieve token and if the " | |
| "probability of Retrieve=Yes divided by the sum of Retrieve=Yes and Retrieve=No " | |
| "exceeds the threshold delta, retrieval is triggered. The default threshold is " | |
| "0.2 for most tasks." | |
| ), | |
| }, | |
| { | |
| "id": "doc_07", | |
| "title": "Segment-Level Beam Search", | |
| "text": ( | |
| "SELF-RAG uses segment-level beam search with beam width B=2 during inference. " | |
| "For each retrieved passage, the model generates a candidate continuation scored " | |
| "using a weighted sum of IsRel, IsSupport, and IsUse token probabilities " | |
| "combined with the log-probability of the generated segment. Default weights " | |
| "are 1.0 for IsRel, 1.0 for IsSupport, and 0.5 for IsUse." | |
| ), | |
| }, | |
| ] | |
| EXAMPLES = [ | |
| ["What is SELF-RAG and how does it use reflection tokens?"], | |
| ["What is the main limitation of Vanilla RAG?"], | |
| ["How does RAG improve factual correctness?"], | |
| ["What is hallucination in language models?"], | |
| ["How does CRAG handle irrelevant retrieved documents?"], | |
| [ | |
| "What is the deployment date and version number of the AuthService " | |
| "microservice in the Pegasus project?" | |
| ], | |
| ] | |
| class Passage: | |
| id: str | |
| title: str | |
| text: str | |
| class CritiqueScores: | |
| p_relevant: float = 0.0 | |
| p_irrelevant: float = 0.0 | |
| p_fully: float = 0.0 | |
| p_partially: float = 0.0 | |
| p_no_support: float = 0.0 | |
| p_use_1: float = 0.0 | |
| p_use_2: float = 0.0 | |
| p_use_3: float = 0.0 | |
| p_use_4: float = 0.0 | |
| p_use_5: float = 0.0 | |
| class SegmentResult: | |
| text: str | |
| retrieve_tok: str | |
| passage: Optional[Passage] = None | |
| isrel: Optional[str] = None | |
| issup: Optional[str] = None | |
| isuse: Optional[str] = None | |
| score: float = 0.0 | |
| is_sufficient: bool = True | |
| isrel_score: float = 0.0 | |
| query_coverage: float = 0.0 | |
| class SelfRAGOutput: | |
| query: str | |
| segments: list = field(default_factory=list) | |
| abstained: bool = False | |
| def answer(self) -> str: | |
| return " ".join(s.text for s in self.segments if s.is_sufficient).strip() | |
| class LightweightRetriever: | |
| def __init__(self, device: str = "cpu"): | |
| self.device = device | |
| self.model = SentenceTransformer("all-MiniLM-L6-v2", device=device) | |
| self.passages = [] | |
| self.index = None | |
| def index_documents(self, docs): | |
| self.passages = [Passage(**d) for d in docs] | |
| texts = [f"{p.title} {p.text}".strip() for p in self.passages] | |
| embs = self.model.encode( | |
| texts, convert_to_numpy=True, normalize_embeddings=True | |
| ).astype("float32") | |
| self.index = faiss.IndexFlatIP(embs.shape[1]) | |
| self.index.add(embs) | |
| def retrieve(self, query: str, k: int = K_PASSAGES): | |
| q_emb = self.model.encode( | |
| [query], convert_to_numpy=True, normalize_embeddings=True | |
| ).astype("float32") | |
| _, idxs = self.index.search(q_emb, k) | |
| return [self.passages[i] for i in idxs[0] if i < len(self.passages)] | |
| retriever = None | |
| gen_tokenizer = None | |
| gen_model = None | |
| TOKEN_IDS = {} | |
| MODEL_LOAD_ERROR = None | |
| def score_isrel(cs: CritiqueScores) -> float: | |
| denom = cs.p_relevant + cs.p_irrelevant | |
| return cs.p_relevant / denom if denom > 0 else 0.0 | |
| def score_issup(cs: CritiqueScores) -> float: | |
| denom = cs.p_fully + cs.p_partially + cs.p_no_support | |
| if denom == 0: | |
| return 0.0 | |
| return (cs.p_fully / denom) + 0.5 * (cs.p_partially / denom) | |
| def score_isuse(cs: CritiqueScores) -> float: | |
| probs = { | |
| 1: cs.p_use_1, | |
| 2: cs.p_use_2, | |
| 3: cs.p_use_3, | |
| 4: cs.p_use_4, | |
| 5: cs.p_use_5, | |
| } | |
| denom = sum(probs.values()) | |
| if denom == 0: | |
| return 0.0 | |
| return sum(IsUseToken.WEIGHTS[i] * p for i, p in probs.items()) / denom | |
| def segment_score(log_prob: float, cs: CritiqueScores) -> float: | |
| return log_prob + W_REL * score_isrel(cs) + W_SUP * score_issup(cs) + W_USE * score_isuse(cs) | |
| def should_retrieve(p_yes: float, p_no: float, delta: float = DELTA) -> bool: | |
| denom = p_yes + p_no | |
| return (p_yes / denom) > delta if denom > 0 else False | |
| def setup_models(): | |
| global retriever, gen_tokenizer, gen_model, TOKEN_IDS, MODEL_LOAD_ERROR | |
| if retriever is not None and gen_tokenizer is not None and gen_model is not None: | |
| return | |
| try: | |
| retriever = LightweightRetriever(device="cpu") | |
| retriever.index_documents(DOCUMENTS) | |
| gen_tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT) | |
| if LOAD_IN_4BIT: | |
| from transformers import BitsAndBytesConfig | |
| quant_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_compute_dtype=torch.float16, | |
| ) | |
| gen_model = AutoModelForCausalLM.from_pretrained( | |
| CHECKPOINT, | |
| device_map="auto", | |
| quantization_config=quant_config, | |
| low_cpu_mem_usage=True, | |
| ) | |
| else: | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError( | |
| "CUDA GPU not available. This notebook-faithful demo expects a GPU " | |
| "Space, or you can set LOAD_IN_4BIT=1 for a quantized fallback." | |
| ) | |
| gen_model = AutoModelForCausalLM.from_pretrained( | |
| CHECKPOINT, | |
| torch_dtype=torch.float16, | |
| device_map="auto", | |
| low_cpu_mem_usage=True, | |
| ) | |
| gen_model.eval() | |
| gen_model.generation_config.do_sample = False | |
| gen_model.generation_config.temperature = None | |
| gen_model.generation_config.top_p = None | |
| def get_tok_id(tok_str: str) -> int: | |
| ids = gen_tokenizer.encode(tok_str, add_special_tokens=False) | |
| return ids[0] | |
| TOKEN_IDS = {tok: get_tok_id(tok) for tok in ALL_REFLECTION_TOKENS} | |
| MODEL_LOAD_ERROR = None | |
| except Exception as exc: # pragma: no cover - startup/runtime guard | |
| MODEL_LOAD_ERROR = str(exc) | |
| raise | |
| def _model_device(): | |
| return next(gen_model.parameters()).device | |
| def _encode_input(text: str): | |
| enc = gen_tokenizer(text, return_tensors="pt", truncation=True, max_length=2048) | |
| device = _model_device() | |
| return {k: v.to(device) for k, v in enc.items()} | |
| def _next_token_probs(prompt: str): | |
| enc = _encode_input(prompt) | |
| logits = gen_model(**enc).logits[0, -1, :] | |
| return torch.softmax(logits, dim=-1) | |
| def _generate_segment(prompt: str): | |
| enc = _encode_input(prompt) | |
| out = gen_model.generate( | |
| **enc, | |
| max_new_tokens=MAX_NEW_TOKS, | |
| do_sample=False, | |
| return_dict_in_generate=True, | |
| output_scores=True, | |
| ) | |
| gen_ids = out.sequences[0, enc["input_ids"].shape[1]:] | |
| text = gen_tokenizer.decode(gen_ids, skip_special_tokens=False).strip() | |
| for tok in ALL_REFLECTION_TOKENS: | |
| if tok in text: | |
| text = text[:text.index(tok)].strip() | |
| log_prob = 0.0 | |
| if out.scores: | |
| for tid, score_t in zip(gen_ids, out.scores): | |
| log_prob += torch.log_softmax(score_t[0], dim=-1)[tid].item() | |
| log_prob /= max(len(out.scores), 1) | |
| return text, log_prob | |
| def _predict_retrieve(prompt: str): | |
| probs = _next_token_probs(prompt) | |
| p_yes = probs[TOKEN_IDS[RetrieveToken.YES]].item() | |
| p_no = probs[TOKEN_IDS[RetrieveToken.NO]].item() | |
| return should_retrieve(p_yes, p_no, DELTA), p_yes, p_no | |
| def _score_passage_relevance(query: str, passage: Passage) -> float: | |
| prompt = f"### Instruction:\n{query}\n\n### Input:\n<p>{passage.text}</p>" | |
| probs = _next_token_probs(prompt) | |
| p_rel = probs[TOKEN_IDS[IsRelToken.RELEVANT]].item() | |
| p_irr = probs[TOKEN_IDS[IsRelToken.IRRELEVANT]].item() | |
| denom = p_rel + p_irr | |
| return p_rel / denom if denom > 0 else 0.0 | |
| def _content_terms(text: str): | |
| stop = { | |
| "a", "an", "the", "is", "are", "was", "were", "be", "been", "being", | |
| "do", "does", "did", "have", "has", "had", "how", "what", "when", | |
| "where", "why", "which", "who", "whom", "this", "that", "these", | |
| "those", "and", "or", "but", "for", "with", "into", "from", "about", | |
| "main", "use", "uses", "using", "used", "number", "version", "date", | |
| "deployment", | |
| } | |
| tokens = re.findall(r"[A-Za-z][A-Za-z0-9_-]+", text.lower()) | |
| return {tok for tok in tokens if tok not in stop and len(tok) > 2} | |
| def _query_passage_coverage(query: str, passage: Passage) -> float: | |
| q_terms = _content_terms(query) | |
| if not q_terms: | |
| return 1.0 | |
| p_terms = _content_terms(f"{passage.title} {passage.text}") | |
| return len(q_terms & p_terms) / len(q_terms) | |
| def _first_segment_abstention_probe(query: str, context: str = ""): | |
| passages = retriever.retrieve(query + (" " + context if context else ""), k=K_PASSAGES) | |
| best_passage = None | |
| best_isrel = -1.0 | |
| best_coverage = 0.0 | |
| for passage in passages: | |
| isrel_score = _score_passage_relevance(query, passage) | |
| coverage = _query_passage_coverage(query, passage) | |
| if (isrel_score, coverage) > (best_isrel, best_coverage): | |
| best_passage = passage | |
| best_isrel = isrel_score | |
| best_coverage = coverage | |
| return passages, best_passage, max(best_isrel, 0.0), best_coverage | |
| def _get_critique_scores(base_prompt: str, passage: Passage, gen_text: str): | |
| cs = CritiqueScores() | |
| isrel_prompt = f"{base_prompt}<p>{passage.text}</p>" | |
| probs = _next_token_probs(isrel_prompt) | |
| cs.p_relevant = probs[TOKEN_IDS[IsRelToken.RELEVANT]].item() | |
| cs.p_irrelevant = probs[TOKEN_IDS[IsRelToken.IRRELEVANT]].item() | |
| issup_prompt = f"{isrel_prompt}{IsRelToken.RELEVANT} {gen_text}" | |
| probs = _next_token_probs(issup_prompt) | |
| cs.p_fully = probs[TOKEN_IDS[IsSupportToken.FULLY]].item() | |
| cs.p_partially = probs[TOKEN_IDS[IsSupportToken.PARTIALLY]].item() | |
| cs.p_no_support = probs[TOKEN_IDS[IsSupportToken.NO]].item() | |
| isuse_prompt = f"{issup_prompt}{IsSupportToken.FULLY}" | |
| probs = _next_token_probs(isuse_prompt) | |
| cs.p_use_5 = probs[TOKEN_IDS[IsUseToken.FIVE]].item() | |
| cs.p_use_4 = probs[TOKEN_IDS[IsUseToken.FOUR]].item() | |
| cs.p_use_3 = probs[TOKEN_IDS[IsUseToken.THREE]].item() | |
| cs.p_use_2 = probs[TOKEN_IDS[IsUseToken.TWO]].item() | |
| cs.p_use_1 = probs[TOKEN_IDS[IsUseToken.ONE]].item() | |
| return cs | |
| def _get_isuse_only(prompt: str, text: str): | |
| cs = CritiqueScores() | |
| probs = _next_token_probs(f"{prompt}{text}") | |
| cs.p_use_5 = probs[TOKEN_IDS[IsUseToken.FIVE]].item() | |
| cs.p_use_4 = probs[TOKEN_IDS[IsUseToken.FOUR]].item() | |
| cs.p_use_3 = probs[TOKEN_IDS[IsUseToken.THREE]].item() | |
| cs.p_use_2 = probs[TOKEN_IDS[IsUseToken.TWO]].item() | |
| cs.p_use_1 = probs[TOKEN_IDS[IsUseToken.ONE]].item() | |
| return cs | |
| def _argmax_tok(options, probs_map): | |
| return max(options, key=lambda t: probs_map.get(t, -1.0)) | |
| def run_selfrag(query: str, max_segments: int = MAX_SEGMENTS) -> SelfRAGOutput: | |
| output = SelfRAGOutput(query=query) | |
| context = "" | |
| for seg_idx in range(max_segments): | |
| retrieve_prompt = ( | |
| f"### Instruction:\n{query}\n\n" | |
| f"### Response:\n{context}" | |
| ) | |
| if seg_idx == 0: | |
| _, probe_passage, probe_isrel, probe_coverage = _first_segment_abstention_probe(query, context) | |
| if ( | |
| probe_passage is None | |
| or probe_coverage < ABSTENTION_MIN_QUERY_COVERAGE | |
| or probe_isrel < ABSTENTION_ISREL_THRESHOLD | |
| ): | |
| output.abstained = True | |
| output.segments.append( | |
| SegmentResult( | |
| text="Insufficient evidence to answer this question.", | |
| retrieve_tok=RetrieveToken.NO, | |
| passage=probe_passage, | |
| isrel=( | |
| IsRelToken.RELEVANT | |
| if probe_isrel >= 0.5 | |
| else IsRelToken.IRRELEVANT | |
| ), | |
| is_sufficient=False, | |
| score=-float("inf"), | |
| isrel_score=max(probe_isrel, 0.0), | |
| query_coverage=probe_coverage, | |
| ) | |
| ) | |
| break | |
| do_retrieve, _, _ = _predict_retrieve(retrieve_prompt) | |
| if do_retrieve: | |
| passages = retriever.retrieve(query + " " + context, k=K_PASSAGES) | |
| candidates = [] | |
| for passage in passages: | |
| gen_prompt = ( | |
| f"### Instruction:\n{query}\n\n" | |
| f"### Input:\n<p>{passage.text}</p>\n\n" | |
| f"### Response:\n{context}" | |
| ) | |
| text, log_prob = _generate_segment(gen_prompt) | |
| if not text: | |
| continue | |
| cs = _get_critique_scores( | |
| base_prompt=f"### Instruction:\n{query}\n\n### Input:\n", | |
| passage=passage, | |
| gen_text=text, | |
| ) | |
| isrel_raw = score_isrel(cs) | |
| isrel_label = ( | |
| IsRelToken.RELEVANT if isrel_raw > 0.5 else IsRelToken.IRRELEVANT | |
| ) | |
| issup_label = ( | |
| IsSupportToken.FULLY | |
| if cs.p_fully >= max(cs.p_partially, cs.p_no_support) | |
| else IsSupportToken.PARTIALLY | |
| if cs.p_partially >= cs.p_no_support | |
| else IsSupportToken.NO | |
| ) | |
| isuse_label = _argmax_tok( | |
| IsUseToken.ALL, | |
| { | |
| IsUseToken.FIVE: cs.p_use_5, | |
| IsUseToken.FOUR: cs.p_use_4, | |
| IsUseToken.THREE: cs.p_use_3, | |
| IsUseToken.TWO: cs.p_use_2, | |
| IsUseToken.ONE: cs.p_use_1, | |
| }, | |
| ) | |
| candidates.append( | |
| SegmentResult( | |
| text=text, | |
| retrieve_tok=RetrieveToken.YES, | |
| passage=passage, | |
| isrel=isrel_label, | |
| issup=issup_label, | |
| isuse=isuse_label, | |
| score=segment_score(log_prob, cs), | |
| is_sufficient=(issup_label != IsSupportToken.NO), | |
| isrel_score=isrel_raw, | |
| query_coverage=_query_passage_coverage(query, passage), | |
| ) | |
| ) | |
| if not candidates: | |
| output.abstained = True | |
| segment = SegmentResult( | |
| text="Insufficient evidence to answer this question.", | |
| retrieve_tok=RetrieveToken.YES, | |
| is_sufficient=False, | |
| score=-float("inf"), | |
| isrel_score=0.0, | |
| ) | |
| else: | |
| candidates.sort(key=lambda candidate: candidate.score, reverse=True) | |
| segment = candidates[0] | |
| else: | |
| text, log_prob = _generate_segment(retrieve_prompt) | |
| cs = _get_isuse_only(retrieve_prompt, text) | |
| isuse_label = _argmax_tok( | |
| IsUseToken.ALL, | |
| { | |
| IsUseToken.FIVE: cs.p_use_5, | |
| IsUseToken.FOUR: cs.p_use_4, | |
| IsUseToken.THREE: cs.p_use_3, | |
| IsUseToken.TWO: cs.p_use_2, | |
| IsUseToken.ONE: cs.p_use_1, | |
| }, | |
| ) | |
| segment = SegmentResult( | |
| text=text, | |
| retrieve_tok=RetrieveToken.NO, | |
| isuse=isuse_label, | |
| score=log_prob + W_USE * score_isuse(cs), | |
| isrel_score=0.0, | |
| ) | |
| output.segments.append(segment) | |
| context = (context + " " + segment.text).strip() | |
| if output.abstained or not segment.text or "</s>" in segment.text: | |
| break | |
| return output | |
| def _fix_spaced_acronyms(text: str) -> str: | |
| pattern = re.compile(r"\b(?:[A-Z]\s+){2,}[A-Z]\b") | |
| while True: | |
| updated = pattern.sub(lambda match: match.group(0).replace(" ", ""), text) | |
| if updated == text: | |
| return updated | |
| text = updated | |
| def _detokenize_text(text: str) -> str: | |
| text = text.replace("\u2581", " ").replace("_", " ") | |
| text = _fix_spaced_acronyms(text) | |
| text = re.sub(r"\s+([,.;:!?])", r"\1", text) | |
| text = re.sub(r"\s*([()\[\]{}])\s*", r" \1 ", text) | |
| text = re.sub(r"\s*-\s*", "-", text) | |
| return " ".join(text.split()).strip() | |
| def clean_answer(text: str) -> str: | |
| for tok in ALL_REFLECTION_TOKENS: | |
| text = text.replace(tok, "") | |
| text = re.sub(r"<[^>]+>", "", text) | |
| text = text.replace("</s>", "").replace("<s>", "") | |
| text = re.sub(r"\[.*?\]", "", text) | |
| text = re.sub(r"^\s*\d+[\.\)]\s*", "", text) | |
| text = re.sub( | |
| r"^(great question!?|sure!?|here's.*?:)\s*", | |
| "", | |
| text, | |
| flags=re.IGNORECASE, | |
| ) | |
| return _detokenize_text(text) | |
| def _norm(text: str) -> str: | |
| text = _detokenize_text(text).lower() | |
| text = re.sub(r"\b(a|an|the)\b", " ", text) | |
| text = text.translate(str.maketrans("", "", string.punctuation)) | |
| return " ".join(text.split()) | |
| def token_f1(pred: str, gold: str) -> float: | |
| pred_tokens = _norm(pred).split() | |
| gold_tokens = _norm(gold).split() | |
| common = Counter(pred_tokens) & Counter(gold_tokens) | |
| overlap = sum(common.values()) | |
| if overlap == 0: | |
| return 0.0 | |
| precision = overlap / len(pred_tokens) if pred_tokens else 0.0 | |
| recall = overlap / len(gold_tokens) if gold_tokens else 0.0 | |
| return 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0.0 | |
| def _split_sentences(text: str): | |
| return [part.strip() for part in re.split(r"(?<=[.!?])\s+", text) if part.strip()] | |
| def _best_source_snippet(question: str, source_text: str) -> str: | |
| sentences = _split_sentences(source_text) | |
| if not sentences: | |
| return source_text.strip() | |
| q_terms = _content_terms(question) | |
| ranked = [] | |
| for sent in sentences: | |
| sent_terms = _content_terms(sent) | |
| overlap = len(q_terms & sent_terms) | |
| ranked.append((overlap, len(sent_terms), sent)) | |
| ranked.sort(reverse=True) | |
| return ranked[0][2].strip() | |
| def _primary_segment(output: SelfRAGOutput) -> Optional[SegmentResult]: | |
| candidates = [seg for seg in output.segments if seg.passage] | |
| if not candidates: | |
| return None | |
| sufficient = [seg for seg in candidates if seg.is_sufficient] | |
| chosen = sufficient or candidates | |
| return max(chosen, key=lambda seg: (seg.score, seg.isrel_score, seg.query_coverage)) | |
| def _compact_answer(output: SelfRAGOutput, source_text: str = "") -> str: | |
| if output.abstained: | |
| return "I don't have enough evidence in the indexed knowledge base to answer that reliably." | |
| cleaned = clean_answer(output.answer) | |
| if not cleaned: | |
| cleaned = source_text | |
| if source_text and (len(cleaned) > 260 or "[" in cleaned or "\n" in cleaned): | |
| cleaned = _best_source_snippet(output.query, source_text) | |
| sentences = _split_sentences(cleaned) | |
| compact = " ".join(sentences[:2]) if sentences else cleaned.strip() | |
| if len(compact) > 260: | |
| compact = compact[:257].rsplit(" ", 1)[0] + "..." | |
| return compact | |
| def _result_card(title: str, body: str, tone: str = "normal") -> str: | |
| palette = { | |
| "normal": ("#ffffff", "#0f172a", "#dbe4f0"), | |
| "error": ("#fff4f4", "#9f1239", "#fecdd3"), | |
| "warn": ("#fff8eb", "#9a3412", "#fed7aa"), | |
| } | |
| bg, fg, border = palette[tone] | |
| return ( | |
| f"<div style='background:{bg};color:{fg};border:1px solid {border};" | |
| "border-radius:18px;padding:18px 20px;box-shadow:0 10px 30px rgba(15,23,42,0.06)'>" | |
| f"<div style='font-size:12px;letter-spacing:0.08em;text-transform:uppercase;" | |
| "opacity:0.65;margin-bottom:8px'>" | |
| f"{title}</div><div style='font-size:16px;line-height:1.65'>{body}</div></div>" | |
| ) | |
| def render_output(output: SelfRAGOutput) -> str: | |
| source_segment = _primary_segment(output) | |
| source_text = source_segment.passage.text if source_segment and source_segment.passage else "" | |
| final_answer = _compact_answer(output, source_text=source_text) | |
| source_html = _result_card( | |
| "Source", | |
| "No supporting passage was selected." | |
| if not source_segment or not source_segment.passage | |
| else ( | |
| f"<div style='font-weight:700;font-size:18px;margin-bottom:8px'>" | |
| f"{source_segment.passage.title}</div>" | |
| f"<div style='font-size:13px;opacity:0.7;margin-bottom:10px'>" | |
| f"{source_segment.passage.id}</div>" | |
| f"<div>{_best_source_snippet(output.query, source_segment.passage.text)}</div>" | |
| ), | |
| tone="warn" if output.abstained else "normal", | |
| ) | |
| answer_html = _result_card( | |
| "Final Answer", | |
| final_answer, | |
| tone="warn" if output.abstained else "normal", | |
| ) | |
| question_html = ( | |
| "<div style='margin-bottom:14px;padding:14px 16px;border-radius:16px;" | |
| "background:linear-gradient(135deg,#eef4ff 0%,#f8fafc 100%);" | |
| "border:1px solid #d9e6ff;color:#0f172a'>" | |
| f"<div style='font-size:12px;letter-spacing:0.08em;text-transform:uppercase;" | |
| "opacity:0.65;margin-bottom:6px'>Question</div>" | |
| f"<div style='font-size:18px;font-weight:600;line-height:1.5'>{output.query}</div>" | |
| "</div>" | |
| ) | |
| return ( | |
| "<div style='font-family:ui-sans-serif,system-ui,sans-serif'>" | |
| f"{question_html}" | |
| "<div style='display:grid;grid-template-columns:1fr;gap:14px'>" | |
| f"{answer_html}{source_html}" | |
| "</div></div>" | |
| ) | |
| def predict(query: str): | |
| query = (query or "").strip() | |
| if not query: | |
| return _result_card( | |
| "Ready", | |
| "Enter a question to see a grounded answer with its source passage.", | |
| ) | |
| try: | |
| setup_models() | |
| output = run_selfrag(query, max_segments=MAX_SEGMENTS) | |
| return render_output(output) | |
| except Exception as exc: # pragma: no cover - runtime guard | |
| global MODEL_LOAD_ERROR | |
| MODEL_LOAD_ERROR = str(exc) | |
| return _result_card( | |
| "Inference Failed", | |
| f"{MODEL_LOAD_ERROR}", | |
| tone="error", | |
| ) | |
| with gr.Blocks( | |
| css=""" | |
| .gradio-container { | |
| background: | |
| radial-gradient(circle at top left, rgba(59,130,246,0.08), transparent 28%), | |
| radial-gradient(circle at top right, rgba(16,185,129,0.08), transparent 22%), | |
| linear-gradient(180deg, #f8fbff 0%, #f4f7fb 100%); | |
| } | |
| #ask_btn { | |
| background: linear-gradient(135deg, #0f172a 0%, #1d4ed8 100%) !important; | |
| color: white !important; | |
| border: none !important; | |
| } | |
| """ | |
| ) as demo: | |
| gr.Markdown( | |
| f""" | |
| # {TOOL_NAME} | |
| {TOOL_TAGLINE} | |
| """ | |
| ) | |
| with gr.Row(): | |
| qbox = gr.Textbox(label="Your Question", scale=5, lines=2) | |
| ask = gr.Button("Ask", elem_id="ask_btn", scale=1) | |
| result = gr.HTML(_result_card("Ready", "Enter a question to see a grounded answer with its source passage.")) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=qbox, | |
| outputs=result, | |
| fn=predict, | |
| cache_examples=False, | |
| ) | |
| ask.click(predict, inputs=qbox, outputs=result) | |
| qbox.submit(predict, inputs=qbox, outputs=result) | |
| if __name__ == "__main__": | |
| demo.queue().launch() | |