""" generator.py -- Dual-backend generation for the RAG-lite QA system. - FLAN-T5: 1. Explicit rules for full sentences, context usage, and question-type-aware instructions. 2. max_new_tokens raised to 220 (from 150) for more complete answers. 3. no_repeat_ngram_size=3 retained to block repetition loops. 4. repetition_penalty=1.3 retained. 5. do_sample=True with temperature=0.3 added -- slight randomness improves reasoning vs pure greedy decoding which always picks the shortest high-probability token sequence. 6. Answer expansion layer: if FLAN returns fewer than 8 words, the system reruns generation with an explicit expansion prompt to produce a full sentence rather than a fragment. 7. Post-processing: capitalize first letter, strip trailing whitespace. 8. Context passed to FLAN is capped at top 3 chunks (not 7) to reduce noise. FLAN is a small model and performs worse with long noisy inputs. Groq improvements: 9. "Answer ONLY the given question" retained to prevent over-answering. 10. max_tokens=512 retained. 11. temperature=0.1 retained. Output cleaning: 12. clean_answer() strips separator characters (===, ---) from any generated output before returning it to app.py. """ import os import re import torch from dotenv import load_dotenv load_dotenv() LOCAL_MODEL_NAME = "google/flan-t5-base" GROQ_MODEL_NAME = os.environ.get("GROQ_MODEL", "llama-3.1-8b-instant") # Cap context at top 3 chunks for FLAN. More chunks add noise that confuses # the small model. Groq handles longer context fine so it uses all chunks. FLAN_MAX_CHUNKS = 3 _local_pipeline = None _FAILURE_PHRASES = {"not found", "not mentioned", "not stated", "unknown", "n/a", "none"} _HEDGING_PHRASES = [ "i cannot", "i don't know", "not available", "not provided", "cannot be found", "not mentioned", "not stated", "no information", ] _LIST_TRIGGER_WORDS = { "types", "type", "kinds", "kind", "categories", "category", "applications", "application", "uses", "use", "examples", "example", "techniques", "technique", "methods", "method", "components", "challenges", "challenge", "features", "feature", "benefits", "advantages", "disadvantages", "steps", "stages", } _COMPARE_TRIGGER_WORDS = { "difference", "different", "differs", "compare", "comparison", "versus", "vs", "unlike", "contrast", "distinguish", } # GROUNDING INSTRUCTION GROUNDING_INSTRUCTION = ( "You are a precise and detail-oriented question-answering assistant. " "Answer the question using only the information in the provided context. " "Do not use any external knowledge or information from your training data. " "You may combine information from multiple sentences within the context " "to form a complete answer. " "When asked about types, techniques, categories, applications, uses, or examples, " "list all items explicitly mentioned in the context. Do not stop after one item. " "Always write complete, grammatically correct sentences. " "Do not copy incomplete phrases. Expand short phrases into full explanations. " "Only say the answer is not found if the information is genuinely absent. " "Do not add facts, inferences, or assumptions beyond what the context states." ) # OUTPUT CLEANING def clean_answer(text: str) -> str: """ Remove separator characters and normalize whitespace. Handles both the === / --- dividers from format_multi_answer_output and any stray markup the model might generate. Also capitalizes the first letter and strips trailing whitespace. """ if not text: return text # Remove lines that are purely separator characters lines = text.split("\n") cleaned_lines = [] for line in lines: stripped = line.strip() # Drop lines made entirely of = or - characters (dividers) if stripped and re.fullmatch(r"[=\-]{3,}", stripped): continue cleaned_lines.append(line) text = "\n".join(cleaned_lines) # Collapse 3+ consecutive newlines to 2 text = re.sub(r"\n{3,}", "\n\n", text) text = text.strip() # Capitalize first character if text and text[0].islower(): text = text[0].upper() + text[1:] return text # LOCAL MODEL LOADER def _get_local_pipeline(): global _local_pipeline if _local_pipeline is None: print(f"[INFO] Loading local model: {LOCAL_MODEL_NAME}") from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline device = 0 if torch.cuda.is_available() else -1 tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_NAME) model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_NAME) _local_pipeline = pipeline( "text2text-generation", model=model, tokenizer=tokenizer, device=device, ) print(f"[INFO] Local model loaded on {'GPU' if device == 0 else 'CPU'}") return _local_pipeline # CHUNK MERGING def merge_chunks(chunks: list[dict], max_chunks: int = None) -> str: """ Merge retrieved chunks into one continuous context string. If max_chunks is set, only the first N chunks (by similarity rank, which is the order they arrive in) are used. """ if not chunks: return "" working = chunks[:max_chunks] if max_chunks else chunks if len(working) == 1: return working[0]["text"].strip() ordered = sorted(working, key=lambda c: c.get("start_word", 0)) merged_words = ordered[0]["text"].split() for chunk in ordered[1:]: chunk_words = chunk["text"].split() overlap_len = 0 max_check = min(len(merged_words), len(chunk_words), 40) for n in range(max_check, 0, -1): if merged_words[-n:] == chunk_words[:n]: overlap_len = n break new_words = chunk_words[overlap_len:] if new_words: merged_words.extend(new_words) return " ".join(merged_words) # ------------------------------------------------------------------ # CONFIDENCE SCORING # ------------------------------------------------------------------ def _score_local_confidence(answer: str, context: str) -> float: """Heuristic confidence for FLAN-T5 based on length and context overlap.""" if not answer or answer.lower().strip() in _FAILURE_PHRASES: return 0.0 words = answer.lower().split() word_count = len(words) context_words = set(context.lower().split()) length_score = min(word_count / 20.0, 1.0) overlap = sum(1 for w in words if w in context_words) overlap_score = overlap / word_count if word_count > 0 else 0.0 return round(min(0.6 * length_score + 0.4 * overlap_score, 1.0), 3) def _score_groq_confidence(answer: str) -> float: """Heuristic confidence for Groq based on hedging, length, sentence count.""" if not answer: return 0.0 answer_lower = answer.lower() hedge_penalty = sum(0.25 for p in _HEDGING_PHRASES if p in answer_lower) hedge_score = max(0.0, 1.0 - hedge_penalty) length_score = min(len(answer.split()) / 30.0, 1.0) sent_count = answer.count(".") + answer.count("?") + answer.count("!") sentence_score = min(sent_count / 3.0, 1.0) return round(min(0.5 * hedge_score + 0.3 * length_score + 0.2 * sentence_score, 1.0), 3) # ------------------------------------------------------------------ # PROMPT BUILDERS # ------------------------------------------------------------------ def _question_type_instruction(question: str) -> str: """ Return a targeted instruction line based on what kind of answer is expected. Three types handled: - List/enumeration: "List all items clearly, separated by commas." - Comparison: "Clearly explain both concepts and highlight the difference." - Default: "Provide a clear and complete answer in 1-2 full sentences." """ q_lower = question.lower() words = [w.rstrip("?.,;:!") for w in q_lower.split()] if any(w in _LIST_TRIGGER_WORDS for w in words): return "List all relevant items explicitly mentioned in the context, separated by commas." if any(w in _COMPARE_TRIGGER_WORDS for w in words): return "Clearly explain both concepts and highlight how they differ, using the context." return "Provide a clear and complete answer in 1-2 full sentences using the context." def _build_local_prompt(context: str, question: str) -> str: """ Prompt for flan-t5-base. Design decisions: - Context comes FIRST so FLAN sees it before the instructions. Small seq2seq models attend more strongly to early tokens. - Instructions are short, numbered, and concrete. Abstract rules like "be thorough" do not work on FLAN. Specific rules like "do not give fragment answers" do. - No numbered format example (removed -- caused the 1...2...3 loop). - Question-type instruction at the end as the final directive before Answer:. """ type_instruction = _question_type_instruction(question) return ( f"Context:\n{context.strip()}\n\n" f"Instructions:\n" f"You are a precise and detail-oriented question-answering assistant.\n" f"Rules:\n" f"- Answer ONLY using the provided context above.\n" f"- Always write a complete, meaningful sentence. Do NOT give fragment answers.\n" f"- Do NOT copy incomplete phrases. Expand them into full explanations.\n" f"- Include important details from the context.\n" f"- Do NOT repeat the question in your answer.\n" f"- Do NOT add external knowledge or hallucinate.\n" f"- Use as much relevant information from the context as possible.\n" f"- {type_instruction}\n\n" f"Question: {question.strip()}\n\n" f"Answer:" ) def _build_expansion_prompt(short_answer: str, context: str, question: str) -> str: """ Expansion prompt used when FLAN returns a very short answer (under 8 words). Asks the model to expand the fragment into a complete sentence. """ return ( f"Context:\n{context.strip()}\n\n" f"Short answer: {short_answer.strip()}\n\n" f"Task: Expand the short answer above into a complete, clear sentence " f"using information from the context. " f"The question was: {question.strip()}\n\n" f"Do NOT add external knowledge. Use only the context.\n\n" f"Complete answer:" ) def _build_groq_messages(context: str, question: str) -> tuple[str, str]: """System and user messages for the Groq API.""" type_instruction = _question_type_instruction(question) user = ( f"Context:\n{context.strip()}\n\n" f"Question: {question.strip()}\n\n" f"Important instructions:\n" f"- Answer ONLY the given question. Do not include information for other questions.\n" f"- {type_instruction}\n" f"- Write complete sentences. Do not give fragment answers." ) return GROUNDING_INSTRUCTION, user # ------------------------------------------------------------------ # GENERATION: LOCAL FLAN-T5 # ------------------------------------------------------------------ def _run_flan(prompt: str) -> str: """ Run a single FLAN-T5 generation call and return the raw output string. Centralizes all generation parameters in one place. """ gen = _get_local_pipeline() result = gen( prompt, max_new_tokens=220, do_sample=True, temperature=0.3, repetition_penalty=1.3, no_repeat_ngram_size=3, truncation=True, ) return result[0]["generated_text"].strip() def _generate_local(context: str, question: str) -> tuple[str, bool, float]: """ Generate answer with flan-t5-base. Steps: 1. Run primary generation with the full structured prompt. 2. If the answer is fewer than 8 words (a fragment), run a second generation pass with an expansion prompt to produce a full sentence. 3. Post-process: clean separators, capitalize first letter. 4. Score confidence. """ print(f"[INFO] Using backend: local ({LOCAL_MODEL_NAME})") try: prompt = _build_local_prompt(context, question) answer = _run_flan(prompt) if not answer or len(answer) < 3: print("[WARN] Local model returned empty answer") return "", False, 0.0 if answer.lower().strip() == question.lower().strip(): print("[WARN] Local model echoed the question") return "", False, 0.0 if answer.lower().strip() in _FAILURE_PHRASES: print(f"[WARN] Local model returned failure phrase: {answer}") return "", False, 0.0 # Expansion pass: if answer is too short, try to expand it if len(answer.split()) < 8: print(f"[INFO] Answer too short ({len(answer.split())} words), running expansion") expansion_prompt = _build_expansion_prompt(answer, context, question) expanded = _run_flan(expansion_prompt) # Only use expanded if it is actually longer and not a failure phrase if (expanded and len(expanded.split()) > len(answer.split()) and expanded.lower().strip() not in _FAILURE_PHRASES): answer = expanded print(f"[INFO] Expansion succeeded: {len(answer.split())} words") # Post-processing: clean and capitalize answer = clean_answer(answer) confidence = _score_local_confidence(answer, context) print(f"[INFO] Local answer | confidence: {confidence:.3f} | words: {len(answer.split())}") return answer, True, confidence except Exception as e: print(f"[ERROR] Local generation failed: {e}") return "", False, 0.0 # ------------------------------------------------------------------ # GENERATION: GROQ (with streaming) # ------------------------------------------------------------------ def stream_groq_answer(context: str, question: str, api_key: str): """ Generator that streams Groq response tokens progressively. Yields the cumulative answer string after each token chunk. """ try: from groq import Groq except ImportError: yield "Groq package not installed. Run: pip install groq" return if not api_key: yield "Groq API key not found. Set GROQ_API_KEY in your .env file." return try: client = Groq(api_key=api_key, timeout=30.0) system_msg, user_msg = _build_groq_messages(context, question) stream = client.chat.completions.create( model=GROQ_MODEL_NAME, messages=[ {"role": "system", "content": system_msg}, {"role": "user", "content": user_msg}, ], temperature=0.1, max_tokens=512, stream=True, ) accumulated = "" for chunk in stream: delta = chunk.choices[0].delta.content if delta: accumulated += delta yield accumulated except Exception as e: err = str(e) print(f"[ERROR] Groq streaming failed: {err}") yield f"Groq error: {err}" def _generate_groq(context: str, question: str, api_key: str) -> tuple[str, bool, float]: """Collect streamed Groq response and return complete answer with confidence.""" print(f"[INFO] Using backend: groq ({GROQ_MODEL_NAME})") full_answer = "" for partial in stream_groq_answer(context, question, api_key): full_answer = partial if not full_answer or len(full_answer) < 5: print("[WARN] Groq returned empty answer") return "", False, 0.0 if any(full_answer.startswith(p) for p in ("Groq error:", "Groq package", "Groq API key")): print(f"[WARN] Groq backend error: {full_answer}") return full_answer, False, 0.0 full_answer = clean_answer(full_answer) confidence = _score_groq_confidence(full_answer) print(f"[INFO] Groq answer | confidence: {confidence:.3f} | words: {len(full_answer.split())}") return full_answer, True, confidence # ------------------------------------------------------------------ # PUBLIC INTERFACE # ------------------------------------------------------------------ def generate_answer( context_chunks: list[dict], question: str, mode: str = "local", groq_api_key: str = "", ) -> tuple[str, bool, float]: """ Main entry point called by app.py after retrieval. For local mode: uses only top FLAN_MAX_CHUNKS (3) to reduce noise. For groq mode: uses all retrieved chunks for maximum context. Returns (answer_text, success_flag, confidence_score). """ if mode == "local": # Limit to top 3 chunks for FLAN -- more chunks add noise merged = merge_chunks(context_chunks, max_chunks=FLAN_MAX_CHUNKS) else: merged = merge_chunks(context_chunks) if not merged.strip(): print("[WARN] merge_chunks returned empty string") return "", False, 0.0 resolved_key = (groq_api_key or "").strip() or os.environ.get("GROQ_API_KEY", "") if mode == "groq": return _generate_groq(merged, question, resolved_key) return _generate_local(merged, question) def fallback_answer(best_chunk: str) -> str: """Return best chunk directly when generation fails entirely.""" print("[WARN] Falling back to raw chunk output") preview = best_chunk.strip()[:500] return f"Most relevant passage from the provided text:\n\n\"{preview}\""