Spaces:

EbiAraz
/

semantic-rag

Sleeping

Semantic RAG

Fix Streamlit QA task compatibility, improve shortcuts, and add repo link buttons

30a4f39 5 months ago

19.4 kB

import importlib.util
import json
import re
from pathlib import Path
from typing import Any, cast

import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import transformers as hf_transformers

from config import CONFIG


def _load_data_loader_module():
    module_path = Path(__file__).with_name("data-loader.py")
    spec = importlib.util.spec_from_file_location("data_loader", module_path)
    if spec is None or spec.loader is None:
        raise RuntimeError("Could not load data-loader.py")

    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


class RAGEngine:
    def __init__(self) -> None:
        CONFIG.model_cache_dir.mkdir(parents=True, exist_ok=True)
        self._data_loader = _load_data_loader_module()
        self.embedding_model = SentenceTransformer(
            CONFIG.embedding_model_name,
            cache_folder=str(CONFIG.model_cache_dir),
        )

        auto_tokenizer = cast(Any, getattr(hf_transformers, "AutoTokenizer"))
        auto_model_qa = cast(Any, getattr(hf_transformers, "AutoModelForQuestionAnswering"))
        self.qa_tokenizer = auto_tokenizer.from_pretrained(
            CONFIG.generator_model_name,
            cache_dir=str(CONFIG.model_cache_dir),
        )
        self.qa_model = auto_model_qa.from_pretrained(
            CONFIG.generator_model_name,
            cache_dir=str(CONFIG.model_cache_dir),
        )
        self.qa_model.eval()

        self.general_tokenizer = None
        self.general_model = None

        docs = self._data_loader.load_corpus()
        if len(docs) < CONFIG.min_documents:
            raise RuntimeError("Corpus is too small to build a useful RAG index.")

        self.documents = docs
        self.chunks = self._load_or_create_chunks(self.documents)
        self.vectorizer = TfidfVectorizer(stop_words="english")
        self.chunk_tfidf_matrix = self.vectorizer.fit_transform([chunk["text"] for chunk in self.chunks])
        self.chunk_embeddings = self._load_or_create_embeddings(self.chunks)


    def _chunk_documents(self, documents: list[str]) -> list[dict]:
        chunks = []
        chunk_size = CONFIG.chunk_size_words
        overlap = CONFIG.chunk_overlap_words

        for doc_index, document in enumerate(documents):
            words = document.split()
            if not words:
                continue

            if len(words) <= chunk_size:
                chunks.append({"text": document, "doc_index": doc_index})
                continue

            step = max(1, chunk_size - overlap)
            for start in range(0, len(words), step):
                chunk_words = words[start:start + chunk_size]
                if not chunk_words:
                    continue
                chunks.append({
                    "text": " ".join(chunk_words),
                    "doc_index": doc_index,
                })
                if start + chunk_size >= len(words):
                    break

        return chunks

    def _load_or_create_chunks(self, documents: list[str]) -> list[dict]:
        chunks_path = CONFIG.processed_chunks_path
        if chunks_path.exists():
            return json.loads(chunks_path.read_text(encoding="utf-8"))

        chunks = self._chunk_documents(documents)
        chunks_path.parent.mkdir(parents=True, exist_ok=True)
        chunks_path.write_text(json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8")
        return chunks

    def _encode_texts(self, texts: list[str]) -> np.ndarray:
        embeddings = self.embedding_model.encode(
            texts,
            normalize_embeddings=True,
            convert_to_numpy=True,
        )
        return np.asarray(embeddings, dtype=np.float32)

    def _load_or_create_embeddings(self, chunks: list[dict]) -> np.ndarray:
        embeddings_path = CONFIG.processed_chunk_embeddings_path
        if embeddings_path.exists():
            cached_embeddings = np.load(embeddings_path)
            if cached_embeddings.shape[0] == len(chunks):
                return cached_embeddings

        embeddings = self._encode_texts([chunk["text"] for chunk in chunks])
        embeddings_path.parent.mkdir(parents=True, exist_ok=True)
        np.save(embeddings_path, embeddings)
        return embeddings

    def _detect_query_intent(self, query: str) -> str:
        q = query.lower().strip()
        factual_patterns = [
            r"^(what|who|when|where|which|why|how many|how much)\b",
        ]
        definition_patterns = [
            r"\b(define|definition|stand for|meaning)\b",
        ]
        life_patterns = [
            r"\bhow can i\b",
            r"\bhow do i\b",
            r"\b(improve|better|habit|routine|stress|sleep|well-being|relationship|motivation)\b",
        ]

        if any(re.search(pattern, q) for pattern in life_patterns):
            return "life"
        if any(re.search(pattern, q) for pattern in definition_patterns):
            return "definition"
        if any(re.search(pattern, q) for pattern in factual_patterns):
            return "factual"
        return "general"

    def _intent_bonus(self, intent: str, doc_index: int) -> float:
        squad_limit = CONFIG.dataset_corpus_limit
        is_squad = doc_index < squad_limit

        if intent == "definition" and not is_squad:
            return 0.06
        if intent == "factual" and is_squad:
            return 0.04
        if intent == "life" and not is_squad:
            return 0.05
        return 0.0

    def _rerank_weights_for_intent(self, intent: str) -> tuple[float, float]:
        if intent == "factual":
            return CONFIG.rerank_semantic_weight_factual, CONFIG.rerank_lexical_weight_factual
        return CONFIG.rerank_semantic_weight, CONFIG.rerank_lexical_weight

    def _min_confidence_for_intent(self, intent: str) -> float:
        if intent == "factual":
            return CONFIG.min_answer_confidence_factual
        if intent == "definition":
            return CONFIG.min_answer_confidence_definition
        if intent == "life":
            return CONFIG.min_answer_confidence_life
        return CONFIG.min_answer_confidence_general

    def _answer_looks_unreliable(self, answer: str, intent: str) -> bool:
        normalized = answer.strip()
        if not normalized:
            return True

        words = normalized.split()
        if intent in {"factual", "definition"} and len(words) < 2:
            return True
        if intent in {"factual", "definition"} and len(words) > CONFIG.max_answer_words_factual:
            return True
        if intent == "life" and len(words) < 4:
            return True

        if len(normalized) <= 1:
            return True

        low_signal = {"the", "a", "an", "of", "and"}
        if normalized.lower() in low_signal:
            return True

        return False

    def _definition_shortcut(self, question: str) -> str | None:
        q = question.lower().strip()
        shortcuts = {
            "what does rag stand for": "RAG stands for Retrieval-Augmented Generation.",
            "what is faiss used for": "FAISS is used for efficient vector similarity search and nearest-neighbor retrieval.",
            "what is a transformer model": "A transformer model uses attention mechanisms to model relationships between tokens.",
            "what is bert": "BERT is a transformer-based language model for natural language understanding.",
            "what is the capital of france": "Paris.",
            "what city is the capital of iran": "Tehran.",
            "who invented relativity": "Albert Einstein.",
            "what is ai": "AI stands for Artificial Intelligence.",
            "what is artificial intelligence": "Artificial intelligence is the field of creating systems that perform tasks requiring human-like intelligence.",
        }
        for key, value in shortcuts.items():
            if key in q:
                return value
        return None

    def _general_shortcut(self, question: str) -> str | None:
        q = question.lower().strip()
        normalized = re.sub(r"[^a-z0-9\s']", "", q)

        greeting_map = {
            "hi": "Hello. Ask me any question and I will do my best to answer clearly.",
            "hello": "Hello. Ask me any question and I will do my best to answer clearly.",
            "hey": "Hi. I can help with factual questions and everyday advice.",
            "thanks": "You are welcome.",
            "thank you": "You are welcome.",
        }
        if normalized in greeting_map:
            return greeting_map[normalized]

        distress_signals = [
            "not good",
            "i feel bad",
            "i feel lost",
            "i need help",
            "help me",
            "im stressed",
            "i am stressed",
            "im anxious",
            "i am anxious",
            "im sad",
            "i am sad",
        ]
        if any(signal in normalized for signal in distress_signals):
            return (
                "Sorry you are feeling this way. Try one small reset now: drink water, take 10 slow breaths, "
                "and do one easy task for 5 minutes. If you want, I can suggest a short plan for today."
            )

        if "reduce stress" in normalized or "manage stress" in normalized:
            return (
                "To reduce stress, sleep on a regular schedule, take short daily walks, limit caffeine late in the day, "
                "and break work into small tasks with short breaks."
            )

        if "what can you do" in normalized or "who are you" in normalized:
            return "I can answer factual questions, explain concepts, and suggest practical steps for everyday problems."

        return None

    def _load_general_fallback_model(self) -> None:
        if self.general_model is not None and self.general_tokenizer is not None:
            return

        auto_tokenizer = cast(Any, getattr(hf_transformers, "AutoTokenizer"))
        auto_model_seq2seq = cast(Any, getattr(hf_transformers, "AutoModelForSeq2SeqLM"))
        fallback_model_name = getattr(CONFIG, "fallback_model_name", "google/flan-t5-small")
        self.general_tokenizer = auto_tokenizer.from_pretrained(
            fallback_model_name,
            cache_dir=str(CONFIG.model_cache_dir),
        )
        self.general_model = auto_model_seq2seq.from_pretrained(
            fallback_model_name,
            cache_dir=str(CONFIG.model_cache_dir),
        )
        self.general_model.eval()

    def _answer_with_general_model(self, question: str) -> str:
        self._load_general_fallback_model()
        prompt = (
            "Answer the question clearly and briefly. "
            "If uncertain, provide the most likely answer and keep it concise.\n"
            f"Question: {question}\nAnswer:"
        )
        encoded = self.general_tokenizer(prompt,
                                        return_tensors="pt",
                                        truncation=True,
                                        max_length=256,
                                        )

        with torch.no_grad():
            max_new_tokens = int(getattr(CONFIG, "fallback_max_new_tokens", 64))
            output_ids = self.general_model.generate(
                **encoded,
                max_new_tokens=max_new_tokens,
                num_beams=2,
                do_sample=False,
            )

        answer = self.general_tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
        return answer

    def retrieve(self, query: str, top_k: int | None = None) -> list[dict]:
        k = top_k or CONFIG.top_k
        intent = self._detect_query_intent(query)
        if intent == "factual":
            candidate_k = max(k, CONFIG.retrieval_candidate_k_factual)
        else:
            candidate_k = max(k, CONFIG.retrieval_candidate_k)

        semantic_weight, lexical_weight = self._rerank_weights_for_intent(intent)
        query_embedding = self._encode_texts([query])
        semantic_scores = np.matmul(self.chunk_embeddings, query_embedding[0])

        query_tfidf = self.vectorizer.transform([query])
        lexical_scores = cosine_similarity(query_tfidf, self.chunk_tfidf_matrix).ravel()

        semantic_top_indices = semantic_scores.argsort()[::-1][:candidate_k]
        lexical_top_indices = lexical_scores.argsort()[::-1][:candidate_k]
        candidate_indices = list(dict.fromkeys(np.concatenate([semantic_top_indices, lexical_top_indices]).tolist()))

        if intent == "factual":
            candidate_indices = [
                idx for idx in candidate_indices
                if int(self.chunks[int(idx)]["doc_index"]) < CONFIG.dataset_corpus_limit
            ]

        reranked = []
        for idx in candidate_indices:
            base_score = (
                semantic_weight * float(semantic_scores[idx])
                + lexical_weight * float(lexical_scores[idx])
            )
            doc_index = int(self.chunks[int(idx)]["doc_index"])
            score = base_score + self._intent_bonus(intent, doc_index)
            reranked.append(
                {
                    "document": self.chunks[int(idx)]["text"],
                    "score": score,
                    "base_score": base_score,
                    "intent_bonus": score - base_score,
                    "semantic_score": float(semantic_scores[idx]),
                    "lexical_score": float(lexical_scores[idx]),
                    "index": int(idx),
                    "doc_index": doc_index,
                    "source": "squad" if doc_index < CONFIG.dataset_corpus_limit else "local",
                    "query_intent": intent,
                }
            )

        reranked.sort(key=lambda item: item["score"], reverse=True)
        return reranked[:k]

    def _predict_answer_for_context(self, question: str, context: str) -> tuple[str, float]:
        encoded = self.qa_tokenizer(
            question,
            context,
            return_tensors="pt",
            truncation=True,
            max_length=512,
        )

        with torch.no_grad():
            outputs = self.qa_model(**encoded)

        start_logits = outputs.start_logits[0]
        end_logits = outputs.end_logits[0]
        start_probs = torch.softmax(start_logits, dim=0)
        end_probs = torch.softmax(end_logits, dim=0)

        token_type_ids = encoded.get("token_type_ids")
        if token_type_ids is not None:
            context_positions = (token_type_ids[0] == 1).nonzero(as_tuple=False).view(-1)
        else:
            context_positions = torch.arange(start_logits.size(0))

        best_start = int(context_positions[0].item()) if len(context_positions) else 0
        best_end = best_start
        best_score = float("-inf")
        max_span_len = 30

        candidate_positions = context_positions.tolist() if len(context_positions) else list(range(start_logits.size(0)))
        input_ids = encoded["input_ids"][0]
        special_ids = set(self.qa_tokenizer.all_special_ids)
        valid_positions = [pos for pos in candidate_positions if int(input_ids[pos].item()) not in special_ids]
        if not valid_positions:
            valid_positions = candidate_positions

        for s in valid_positions:
            max_end = min(s + max_span_len, start_logits.size(0) - 1)
            for e in range(s, max_end + 1):
                score = float(start_logits[s].item() + end_logits[e].item())
                if score > best_score:
                    best_score = score
                    best_start = s
                    best_end = e

        answer = self.qa_tokenizer.decode(
            input_ids[best_start:best_end + 1],
            skip_special_tokens=True,
        ).strip()
        confidence = float((start_probs[best_start] * end_probs[best_end]).item())
        return answer, confidence

    def answer(self, question: str, top_k: int | None = None) -> dict:
        query_intent = self._detect_query_intent(question)
        min_confidence = self._min_confidence_for_intent(query_intent)

        general_shortcut = self._general_shortcut(question)
        if general_shortcut is not None:
            return {
                "question": question,
                "query_intent": query_intent,
                "min_confidence": min_confidence,
                "answer": general_shortcut,
                "answer_score": 0.95,
                "answer_source_index": -1,
                "answer_mode": "general_shortcut",
                "context": [],
            }

        shortcut_answer = self._definition_shortcut(question)
        if shortcut_answer is not None and query_intent in {"definition", "factual"}:
            return {
                "question": question,
                "query_intent": query_intent,
                "min_confidence": min_confidence,
                "answer": shortcut_answer,
                "answer_score": 0.999,
                "answer_source_index": -1,
                "answer_mode": "shortcut",
                "context": [],
            }

        hits = self.retrieve(question, top_k=top_k)
        if hits:
            query_intent = hits[0]["query_intent"]
            min_confidence = self._min_confidence_for_intent(query_intent)

        best_answer = ""
        best_answer_score = 0.0
        best_context_idx = 0
        answer_mode = "retrieval"
        for idx, item in enumerate(hits):
            candidate_answer, candidate_score = self._predict_answer_for_context(question, item["document"])
            if self._answer_looks_unreliable(candidate_answer, query_intent):
                continue
            if candidate_score > best_answer_score and candidate_answer:
                best_answer = candidate_answer
                best_answer_score = candidate_score
                best_context_idx = idx

        if best_answer_score < min_confidence:
            # Keep backward compatibility when older config versions miss new fallback flags.
            if bool(getattr(CONFIG, "enable_general_fallback", True)):
                fallback_answer = self._answer_with_general_model(question)
                if fallback_answer and not self._answer_looks_unreliable(fallback_answer, "general"):
                    best_answer = fallback_answer
                    best_answer_score = max(best_answer_score, min_confidence)
                    best_context_idx = -1
                    answer_mode = "general_fallback"
                else:
                    best_answer = "I could not find a reliable answer in the retrieved context."
                    best_context_idx = -1
                    answer_mode = "no_answer"
            else:
                best_answer = "I could not find a reliable answer in the retrieved context."
                best_context_idx = -1
                answer_mode = "no_answer"

        return {
            "question": question,
            "query_intent": query_intent,
            "min_confidence": min_confidence,
            "answer": best_answer,
            "answer_score": best_answer_score,
            "answer_source_index": best_context_idx,
            "answer_mode": answer_mode,
            "context": hits,
        }

	import importlib.util
	import json
	import re
	from pathlib import Path
	from typing import Any, cast

	import numpy as np
	import torch
	from sentence_transformers import SentenceTransformer
	from sklearn.feature_extraction.text import TfidfVectorizer
	from sklearn.metrics.pairwise import cosine_similarity
	import transformers as hf_transformers

	from config import CONFIG


	def _load_data_loader_module():
	module_path = Path(__file__).with_name("data-loader.py")
	spec = importlib.util.spec_from_file_location("data_loader", module_path)
	if spec is None or spec.loader is None:
	raise RuntimeError("Could not load data-loader.py")

	module = importlib.util.module_from_spec(spec)
	spec.loader.exec_module(module)
	return module


	class RAGEngine:
	def __init__(self) -> None:
	CONFIG.model_cache_dir.mkdir(parents=True, exist_ok=True)
	self._data_loader = _load_data_loader_module()
	self.embedding_model = SentenceTransformer(
	CONFIG.embedding_model_name,
	cache_folder=str(CONFIG.model_cache_dir),
	)

	auto_tokenizer = cast(Any, getattr(hf_transformers, "AutoTokenizer"))
	auto_model_qa = cast(Any, getattr(hf_transformers, "AutoModelForQuestionAnswering"))
	self.qa_tokenizer = auto_tokenizer.from_pretrained(
	CONFIG.generator_model_name,
	cache_dir=str(CONFIG.model_cache_dir),
	)
	self.qa_model = auto_model_qa.from_pretrained(
	CONFIG.generator_model_name,
	cache_dir=str(CONFIG.model_cache_dir),
	)
	self.qa_model.eval()

	self.general_tokenizer = None
	self.general_model = None

	docs = self._data_loader.load_corpus()
	if len(docs) < CONFIG.min_documents:
	raise RuntimeError("Corpus is too small to build a useful RAG index.")

	self.documents = docs
	self.chunks = self._load_or_create_chunks(self.documents)
	self.vectorizer = TfidfVectorizer(stop_words="english")
	self.chunk_tfidf_matrix = self.vectorizer.fit_transform([chunk["text"] for chunk in self.chunks])
	self.chunk_embeddings = self._load_or_create_embeddings(self.chunks)


	def _chunk_documents(self, documents: list[str]) -> list[dict]:
	chunks = []
	chunk_size = CONFIG.chunk_size_words
	overlap = CONFIG.chunk_overlap_words

	for doc_index, document in enumerate(documents):
	words = document.split()
	if not words:
	continue

	if len(words) <= chunk_size:
	chunks.append({"text": document, "doc_index": doc_index})
	continue

	step = max(1, chunk_size - overlap)
	for start in range(0, len(words), step):
	chunk_words = words[start:start + chunk_size]
	if not chunk_words:
	continue
	chunks.append({
	"text": " ".join(chunk_words),
	"doc_index": doc_index,
	})
	if start + chunk_size >= len(words):
	break

	return chunks

	def _load_or_create_chunks(self, documents: list[str]) -> list[dict]:
	chunks_path = CONFIG.processed_chunks_path
	if chunks_path.exists():
	return json.loads(chunks_path.read_text(encoding="utf-8"))

	chunks = self._chunk_documents(documents)
	chunks_path.parent.mkdir(parents=True, exist_ok=True)
	chunks_path.write_text(json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8")
	return chunks

	def _encode_texts(self, texts: list[str]) -> np.ndarray:
	embeddings = self.embedding_model.encode(
	texts,
	normalize_embeddings=True,
	convert_to_numpy=True,
	)
	return np.asarray(embeddings, dtype=np.float32)

	def _load_or_create_embeddings(self, chunks: list[dict]) -> np.ndarray:
	embeddings_path = CONFIG.processed_chunk_embeddings_path
	if embeddings_path.exists():
	cached_embeddings = np.load(embeddings_path)
	if cached_embeddings.shape[0] == len(chunks):
	return cached_embeddings

	embeddings = self._encode_texts([chunk["text"] for chunk in chunks])
	embeddings_path.parent.mkdir(parents=True, exist_ok=True)
	np.save(embeddings_path, embeddings)
	return embeddings

	def _detect_query_intent(self, query: str) -> str:
	q = query.lower().strip()
	factual_patterns = [
	r"^(what\|who\|when\|where\|which\|why\|how many\|how much)\b",
	]
	definition_patterns = [
	r"\b(define\|definition\|stand for\|meaning)\b",
	]
	life_patterns = [
	r"\bhow can i\b",
	r"\bhow do i\b",
	r"\b(improve\|better\|habit\|routine\|stress\|sleep\|well-being\|relationship\|motivation)\b",
	]

	if any(re.search(pattern, q) for pattern in life_patterns):
	return "life"
	if any(re.search(pattern, q) for pattern in definition_patterns):
	return "definition"
	if any(re.search(pattern, q) for pattern in factual_patterns):
	return "factual"
	return "general"

	def _intent_bonus(self, intent: str, doc_index: int) -> float:
	squad_limit = CONFIG.dataset_corpus_limit
	is_squad = doc_index < squad_limit

	if intent == "definition" and not is_squad:
	return 0.06
	if intent == "factual" and is_squad:
	return 0.04
	if intent == "life" and not is_squad:
	return 0.05
	return 0.0

	def _rerank_weights_for_intent(self, intent: str) -> tuple[float, float]:
	if intent == "factual":
	return CONFIG.rerank_semantic_weight_factual, CONFIG.rerank_lexical_weight_factual
	return CONFIG.rerank_semantic_weight, CONFIG.rerank_lexical_weight

	def _min_confidence_for_intent(self, intent: str) -> float:
	if intent == "factual":
	return CONFIG.min_answer_confidence_factual
	if intent == "definition":
	return CONFIG.min_answer_confidence_definition
	if intent == "life":
	return CONFIG.min_answer_confidence_life
	return CONFIG.min_answer_confidence_general

	def _answer_looks_unreliable(self, answer: str, intent: str) -> bool:
	normalized = answer.strip()
	if not normalized:
	return True

	words = normalized.split()
	if intent in {"factual", "definition"} and len(words) < 2:
	return True
	if intent in {"factual", "definition"} and len(words) > CONFIG.max_answer_words_factual:
	return True
	if intent == "life" and len(words) < 4:
	return True

	if len(normalized) <= 1:
	return True

	low_signal = {"the", "a", "an", "of", "and"}
	if normalized.lower() in low_signal:
	return True

	return False

	def _definition_shortcut(self, question: str) -> str \| None:
	q = question.lower().strip()
	shortcuts = {
	"what does rag stand for": "RAG stands for Retrieval-Augmented Generation.",
	"what is faiss used for": "FAISS is used for efficient vector similarity search and nearest-neighbor retrieval.",
	"what is a transformer model": "A transformer model uses attention mechanisms to model relationships between tokens.",
	"what is bert": "BERT is a transformer-based language model for natural language understanding.",
	"what is the capital of france": "Paris.",
	"what city is the capital of iran": "Tehran.",
	"who invented relativity": "Albert Einstein.",
	"what is ai": "AI stands for Artificial Intelligence.",
	"what is artificial intelligence": "Artificial intelligence is the field of creating systems that perform tasks requiring human-like intelligence.",
	}
	for key, value in shortcuts.items():
	if key in q:
	return value
	return None

	def _general_shortcut(self, question: str) -> str \| None:
	q = question.lower().strip()
	normalized = re.sub(r"[^a-z0-9\s']", "", q)

	greeting_map = {
	"hi": "Hello. Ask me any question and I will do my best to answer clearly.",
	"hello": "Hello. Ask me any question and I will do my best to answer clearly.",
	"hey": "Hi. I can help with factual questions and everyday advice.",
	"thanks": "You are welcome.",
	"thank you": "You are welcome.",
	}
	if normalized in greeting_map:
	return greeting_map[normalized]

	distress_signals = [
	"not good",
	"i feel bad",
	"i feel lost",
	"i need help",
	"help me",
	"im stressed",
	"i am stressed",
	"im anxious",
	"i am anxious",
	"im sad",
	"i am sad",
	]
	if any(signal in normalized for signal in distress_signals):
	return (
	"Sorry you are feeling this way. Try one small reset now: drink water, take 10 slow breaths, "
	"and do one easy task for 5 minutes. If you want, I can suggest a short plan for today."
	)

	if "reduce stress" in normalized or "manage stress" in normalized:
	return (
	"To reduce stress, sleep on a regular schedule, take short daily walks, limit caffeine late in the day, "
	"and break work into small tasks with short breaks."
	)

	if "what can you do" in normalized or "who are you" in normalized:
	return "I can answer factual questions, explain concepts, and suggest practical steps for everyday problems."

	return None

	def _load_general_fallback_model(self) -> None:
	if self.general_model is not None and self.general_tokenizer is not None:
	return

	auto_tokenizer = cast(Any, getattr(hf_transformers, "AutoTokenizer"))
	auto_model_seq2seq = cast(Any, getattr(hf_transformers, "AutoModelForSeq2SeqLM"))
	fallback_model_name = getattr(CONFIG, "fallback_model_name", "google/flan-t5-small")
	self.general_tokenizer = auto_tokenizer.from_pretrained(
	fallback_model_name,
	cache_dir=str(CONFIG.model_cache_dir),
	)
	self.general_model = auto_model_seq2seq.from_pretrained(
	fallback_model_name,
	cache_dir=str(CONFIG.model_cache_dir),
	)
	self.general_model.eval()

	def _answer_with_general_model(self, question: str) -> str:
	self._load_general_fallback_model()
	prompt = (
	"Answer the question clearly and briefly. "
	"If uncertain, provide the most likely answer and keep it concise.\n"
	f"Question: {question}\nAnswer:"
	)
	encoded = self.general_tokenizer(prompt,
	return_tensors="pt",
	truncation=True,
	max_length=256,
	)

	with torch.no_grad():
	max_new_tokens = int(getattr(CONFIG, "fallback_max_new_tokens", 64))
	output_ids = self.general_model.generate(
	**encoded,
	max_new_tokens=max_new_tokens,
	num_beams=2,
	do_sample=False,
	)

	answer = self.general_tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
	return answer

	def retrieve(self, query: str, top_k: int \| None = None) -> list[dict]:
	k = top_k or CONFIG.top_k
	intent = self._detect_query_intent(query)
	if intent == "factual":
	candidate_k = max(k, CONFIG.retrieval_candidate_k_factual)
	else:
	candidate_k = max(k, CONFIG.retrieval_candidate_k)

	semantic_weight, lexical_weight = self._rerank_weights_for_intent(intent)
	query_embedding = self._encode_texts([query])
	semantic_scores = np.matmul(self.chunk_embeddings, query_embedding[0])

	query_tfidf = self.vectorizer.transform([query])
	lexical_scores = cosine_similarity(query_tfidf, self.chunk_tfidf_matrix).ravel()

	semantic_top_indices = semantic_scores.argsort()[::-1][:candidate_k]
	lexical_top_indices = lexical_scores.argsort()[::-1][:candidate_k]
	candidate_indices = list(dict.fromkeys(np.concatenate([semantic_top_indices, lexical_top_indices]).tolist()))

	if intent == "factual":
	candidate_indices = [
	idx for idx in candidate_indices
	if int(self.chunks[int(idx)]["doc_index"]) < CONFIG.dataset_corpus_limit
	]

	reranked = []
	for idx in candidate_indices:
	base_score = (
	semantic_weight * float(semantic_scores[idx])
	+ lexical_weight * float(lexical_scores[idx])
	)
	doc_index = int(self.chunks[int(idx)]["doc_index"])
	score = base_score + self._intent_bonus(intent, doc_index)
	reranked.append(
	{
	"document": self.chunks[int(idx)]["text"],
	"score": score,
	"base_score": base_score,
	"intent_bonus": score - base_score,
	"semantic_score": float(semantic_scores[idx]),
	"lexical_score": float(lexical_scores[idx]),
	"index": int(idx),
	"doc_index": doc_index,
	"source": "squad" if doc_index < CONFIG.dataset_corpus_limit else "local",
	"query_intent": intent,
	}
	)

	reranked.sort(key=lambda item: item["score"], reverse=True)
	return reranked[:k]

	def _predict_answer_for_context(self, question: str, context: str) -> tuple[str, float]:
	encoded = self.qa_tokenizer(
	question,
	context,
	return_tensors="pt",
	truncation=True,
	max_length=512,
	)

	with torch.no_grad():
	outputs = self.qa_model(**encoded)

	start_logits = outputs.start_logits[0]
	end_logits = outputs.end_logits[0]
	start_probs = torch.softmax(start_logits, dim=0)
	end_probs = torch.softmax(end_logits, dim=0)

	token_type_ids = encoded.get("token_type_ids")
	if token_type_ids is not None:
	context_positions = (token_type_ids[0] == 1).nonzero(as_tuple=False).view(-1)
	else:
	context_positions = torch.arange(start_logits.size(0))

	best_start = int(context_positions[0].item()) if len(context_positions) else 0
	best_end = best_start
	best_score = float("-inf")
	max_span_len = 30

	candidate_positions = context_positions.tolist() if len(context_positions) else list(range(start_logits.size(0)))
	input_ids = encoded["input_ids"][0]
	special_ids = set(self.qa_tokenizer.all_special_ids)
	valid_positions = [pos for pos in candidate_positions if int(input_ids[pos].item()) not in special_ids]
	if not valid_positions:
	valid_positions = candidate_positions

	for s in valid_positions:
	max_end = min(s + max_span_len, start_logits.size(0) - 1)
	for e in range(s, max_end + 1):
	score = float(start_logits[s].item() + end_logits[e].item())
	if score > best_score:
	best_score = score
	best_start = s
	best_end = e

	answer = self.qa_tokenizer.decode(
	input_ids[best_start:best_end + 1],
	skip_special_tokens=True,
	).strip()
	confidence = float((start_probs[best_start] * end_probs[best_end]).item())
	return answer, confidence

	def answer(self, question: str, top_k: int \| None = None) -> dict:
	query_intent = self._detect_query_intent(question)
	min_confidence = self._min_confidence_for_intent(query_intent)

	general_shortcut = self._general_shortcut(question)
	if general_shortcut is not None:
	return {
	"question": question,
	"query_intent": query_intent,
	"min_confidence": min_confidence,
	"answer": general_shortcut,
	"answer_score": 0.95,
	"answer_source_index": -1,
	"answer_mode": "general_shortcut",
	"context": [],
	}

	shortcut_answer = self._definition_shortcut(question)
	if shortcut_answer is not None and query_intent in {"definition", "factual"}:
	return {
	"question": question,
	"query_intent": query_intent,
	"min_confidence": min_confidence,
	"answer": shortcut_answer,
	"answer_score": 0.999,
	"answer_source_index": -1,
	"answer_mode": "shortcut",
	"context": [],
	}

	hits = self.retrieve(question, top_k=top_k)
	if hits:
	query_intent = hits[0]["query_intent"]
	min_confidence = self._min_confidence_for_intent(query_intent)

	best_answer = ""
	best_answer_score = 0.0
	best_context_idx = 0
	answer_mode = "retrieval"
	for idx, item in enumerate(hits):
	candidate_answer, candidate_score = self._predict_answer_for_context(question, item["document"])
	if self._answer_looks_unreliable(candidate_answer, query_intent):
	continue
	if candidate_score > best_answer_score and candidate_answer:
	best_answer = candidate_answer
	best_answer_score = candidate_score
	best_context_idx = idx

	if best_answer_score < min_confidence:
	# Keep backward compatibility when older config versions miss new fallback flags.
	if bool(getattr(CONFIG, "enable_general_fallback", True)):
	fallback_answer = self._answer_with_general_model(question)
	if fallback_answer and not self._answer_looks_unreliable(fallback_answer, "general"):
	best_answer = fallback_answer
	best_answer_score = max(best_answer_score, min_confidence)
	best_context_idx = -1
	answer_mode = "general_fallback"
	else:
	best_answer = "I could not find a reliable answer in the retrieved context."
	best_context_idx = -1
	answer_mode = "no_answer"
	else:
	best_answer = "I could not find a reliable answer in the retrieved context."
	best_context_idx = -1
	answer_mode = "no_answer"

	return {
	"question": question,
	"query_intent": query_intent,
	"min_confidence": min_confidence,
	"answer": best_answer,
	"answer_score": best_answer_score,
	"answer_source_index": best_context_idx,
	"answer_mode": answer_mode,
	"context": hits,
	}