Update rag/rag_generator.py

rag/rag_generator.py

@@ -29,7 +29,7 @@ from __future__ import annotations
 import os
 import re
 import torch
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+from transformers import BartForConditionalGeneration, BartTokenizer
 
 
 # ------------------------------------------------------------
@@ -38,8 +38,8 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer
 
 MODEL_NAME = "facebook/bart-large"
 
-_tokenizer: T5Tokenizer | None = None
-_model: T5ForConditionalGeneration | None = None
+_tokenizer: BartTokenizer | None = None
+_model: BartForConditionalGeneration | None = None
 
 _MAX_INPUT_TOKENS = 1020
 _DEFAULT_MAX_NEW_TOKENS = 256
@@ -47,15 +47,15 @@ _DEFAULT_MAX_NEW_TOKENS = 256
 _CONTEXT_CHAR_CAP = 2400
 
 
-def _get_model() -> tuple[T5Tokenizer, T5ForConditionalGeneration]:
+def _get_model() -> tuple[BartTokenizer, BartForConditionalGeneration]:
     global _tokenizer, _model
     if _tokenizer is None or _model is None:
-            _tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
-            _model = T5ForConditionalGeneration.from_pretrained(
-                MODEL_NAME,
-                device_map="auto",
-                torch_dtype=torch.float32,
-            )
+        _tokenizer = BartTokenizer.from_pretrained(MODEL_NAME)
+        _model = BartForConditionalGeneration.from_pretrained(
+            MODEL_NAME,
+            device_map="auto",
+            torch_dtype=torch.float32,
+        )
     return _tokenizer, _model
 
 
@@ -136,7 +136,7 @@ def _looks_like_echo_or_garbage(text: str) -> bool:
         return True
 
     if "." not in s and "match" not in low and "conflict" not in low:
-        return True
+            return True
 
     return False
 
@@ -147,12 +147,16 @@ def _looks_like_echo_or_garbage(text: str) -> bool:
 
 _KEY_MATCHES_HEADER_RE = re.compile(r"^\s*KEY MATCHES\s*:\s*$", re.IGNORECASE)
 _CONFLICTS_HEADER_RE = re.compile(r"^\s*CONFLICTS\b.*:\s*$", re.IGNORECASE)
-_CONFLICTS_INLINE_NONE_RE = re.compile(r"^\s*CONFLICTS\s*:\s*not specified\.?\s*$", re.IGNORECASE)
+_CONFLICTS_INLINE_NONE_RE = re.compile(
+    r"^\s*CONFLICTS\s*:\s*not specified\.?\s*$",
+    re.IGNORECASE,
+)
 
 _MATCH_LINE_RE = re.compile(
     r"^\s*-\s*([^:]+)\s*:\s*(.+?)\s*\(matches reference:\s*(.+?)\)\s*$",
     re.IGNORECASE,
 )
+
 _CONFLICT_LINE_RE = re.compile(
     r"^\s*-\s*([^:]+)\s*:\s*(.+?)\s*\(conflicts reference:\s*(.+?)\)\s*$",
     re.IGNORECASE,
@@ -161,7 +165,10 @@ _CONFLICT_LINE_RE = re.compile(
 _GENERIC_BULLET_RE = re.compile(r"^\s*-\s*(.+?)\s*$")
 
 
-def _extract_key_traits_and_conflicts(shaped_ctx: str) -> tuple[list[str], list[str], bool]:
+def _extract_key_traits_and_conflicts(
+    shaped_ctx: str,
+) -> tuple[list[str], list[str], bool]:
+
     key_traits: list[str] = []
     conflicts: list[str] = []
 
@@ -247,7 +254,7 @@ def _format_bullets(items: list[str], *, none_text: str) -> str:
 
 
 # ------------------------------------------------------------
-# NEW — CONFIDENCE STATE EXTRACTOR
+# CONFIDENCE STATE EXTRACTOR
 # ------------------------------------------------------------
 
 _CONF_STATE_RE = re.compile(r"^\s*-\s*Confidence State:\s*(.+?)\s*$", re.IGNORECASE)
@@ -274,7 +281,7 @@ def _extract_confidence_state(shaped_ctx: str) -> tuple[str | None, str | None]:
 
 
 # ------------------------------------------------------------
-# DETERMINISTIC TEMPLATES (PRIMARY AUTHORITY)
+# DETERMINISTIC TEMPLATES
 # ------------------------------------------------------------
 
 def _template_conclusion(
@@ -292,12 +299,10 @@ def _template_conclusion(
     c_short = ", ".join(conflicts[:3]) if conflicts else None
 
     if confidence_state is None:
-        # fallback to legacy deterministic behaviour
         return _deterministic_conclusion(g, key_traits, conflicts)
 
     cs = confidence_state.lower()
 
-    # ---------- STRONG MATCH ----------
     if "strong" in cs:
         return (
             f"This phenotype is indicative of {g} with no conflicting traits observed. "
@@ -305,7 +310,6 @@ def _template_conclusion(
             f"{rec}".strip()
         )
 
-    # ---------- PROBABLE MATCH ----------
     if "probable" in cs or "conflicts present" in cs or "cautious" in cs:
         if c_short:
             return (
@@ -319,26 +323,23 @@ def _template_conclusion(
                 f"reduces certainty. {rec}".strip()
             )
 
-    # ---------- INCONCLUSIVE ----------
     if "inconclusive" in cs or "conflicting" in cs:
         return (
             f"The top genus match is {g}; however, the phenotype is inconclusive due to conflicting "
             f"test results ({c_short or 'multiple conflicting traits'}). {rec}".strip()
         )
 
-    # ---------- WEAK EVIDENCE ----------
     if "weak" in cs:
         return (
             f"The available phenotype provides weak evidence for {g}. "
             f"Additional testing or phenotype data is recommended. {rec}".strip()
         )
 
-    # Unknown confidence state → safe fallback
     return _deterministic_conclusion(g, key_traits, conflicts)
 
 
 # ------------------------------------------------------------
-# LEGACY DETERMINISTIC CONCLUSION (FINAL BACKSTOP)
+# BACKSTOP DETERMINISTIC CONCLUSION
 # ------------------------------------------------------------
 
 def _deterministic_conclusion(genus: str, key_traits: list[str], conflicts: list[str]) -> str:
@@ -405,7 +406,6 @@ def generate_genus_rag_explanation(
             "No reference evidence was available to evaluate this genus against the observed phenotype."
         )
 
-    # Extract evidence sections
     key_traits, conflicts, saw_headers = _extract_key_traits_and_conflicts(context)
 
     if (not saw_headers) or (not key_traits and not conflicts):
@@ -417,10 +417,8 @@ def generate_genus_rag_explanation(
     key_traits_text = _format_bullets(key_traits, none_text="- Not specified.")
     conflicts_text = _format_bullets(conflicts, none_text="- Not specified.")
 
-    # ---------- NEW: read CONFIDENCE STATE ----------
     confidence_state, recommended_action = _extract_confidence_state(context)
 
-    # Generate deterministic template conclusion first
     template_conclusion = _template_conclusion(
         genus_clean,
         confidence_state,
@@ -429,11 +427,10 @@ def generate_genus_rag_explanation(
         recommended_action,
     )
 
-    # If no confidence state → skip LLM entirely
     if confidence_state is None:
         final_conclusion = template_conclusion
+
     else:
-        # Ask LLM to paraphrase but not reinterpret
         prompt = RAG_PROMPT.format(
             genus=genus_clean,
             confidence_state=confidence_state,
@@ -456,7 +453,7 @@ def generate_genus_rag_explanation(
             **inputs,
             max_new_tokens=max_new_tokens,
             temperature=0.0,
-            num_beams=1,
+            num_beams=3,              # BART benefits from small beam search
             do_sample=False,
             repetition_penalty=1.2,
             no_repeat_ngram_size=3,
@@ -469,7 +466,6 @@ def generate_genus_rag_explanation(
             _log_block("RAW OUTPUT (CONCLUSION)", decoded)
             _log_block("CLEANED OUTPUT", cleaned)
 
-        # If LLM garbage → revert to deterministic template
         if _looks_like_echo_or_garbage(cleaned):
             final_conclusion = template_conclusion
             if RAG_GEN_LOG_OUTPUT:
@@ -477,7 +473,6 @@ def generate_genus_rag_explanation(
         else:
             final_conclusion = cleaned
 
-    # ---------- Final structured output ----------
     final = (
         "KEY TRAITS:\n"
         f"{key_traits_text}\n\n"