Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import json
 import pickle
 import re
-from typing import List, Tuple
+from typing import List
 
 import faiss
 import gradio as gr
@@ -23,11 +23,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 # FINAL_TOP_N=4
 # SCORE_THRESHOLD=0.40
 # USE_SCORE_THRESHOLD=true
+# ENABLE_EVIDENCE_SYNTHESIS=true
+# MAX_NEW_TOKENS=500
 #
 # In Space Settings -> Secrets:
 # HF_TOKEN=hf_xxxxxxxxx
 #
-# Optional later:
+# Optional:
 # EMBED_MODEL_NAME=sentence-transformers/all-MiniLM-L6-v2
 
 BASE_MODEL_PATH = os.getenv("BASE_MODEL_PATH", "meta-llama/Llama-3.1-8B-Instruct")
@@ -93,14 +95,15 @@ SYSTEM_PROMPT = (
     "   - For devices/procedures questions: explain the indication, principle, and major cardiac or ECG relevance.\n"
     "   - For special populations questions: adapt the explanation to the relevant group such as pediatrics, pregnancy, older adults, athletes, or patients with congenital/inherited disease.\n"
     "7) When relevant, include hallmark ECG features, diagnostic characteristics, mechanism, clinical implications, major differential points, and management relevance.\n"
-    "8) Keep simple questions concise in 1-3 sentences.\n"
-    "9) For broader or more complex questions, provide a focused multi-sentence explanation that remains concise but sufficiently complete.\n"
-    "10) Do NOT fabricate facts, criteria, thresholds, interpretations, or treatment recommendations.\n"
-    "11) Do NOT include conversational filler, vague generalities, unnecessary repetition, or non-cardiology content.\n"
-    "12) Do NOT ask follow-up questions unless clarification is absolutely required to avoid an unsafe or invalid answer.\n"
-    "13) Use a confident clinical tone.\n"
-    "14) End with a complete declarative sentence.\n"
-    "15) If uncertain, respond exactly with: I’m not sure about that."
+    "8) For simple questions, provide a concise but clinically informative answer.\n"
+    "9) For broader or more complex questions, provide a fuller explanation with clinically relevant detail when supported by the CONTEXT.\n"
+    "10) When the question is general, explanatory, or educational, expand the answer enough to cover the key clinical points.\n"
+    "11) Do NOT fabricate facts, criteria, thresholds, interpretations, or treatment recommendations.\n"
+    "12) Do NOT include conversational filler, vague generalities, unnecessary repetition, or non-cardiology content.\n"
+    "13) Do NOT ask follow-up questions unless clarification is absolutely required to avoid an unsafe or invalid answer.\n"
+    "14) Use a confident clinical tone.\n"
+    "15) End with a complete declarative sentence.\n"
+    "16) If uncertain, respond exactly with: I’m not sure about that."
 )
 
 RAG_INSTRUCTION = (
@@ -112,7 +115,7 @@ RAG_INSTRUCTION = (
     "5) If different CONTEXT passages conflict, prefer the most directly relevant and specific passage.\n"
     "6) If the CONTEXT is relevant, generate a clinically structured answer based on the CONTEXT.\n"
     "7) If the CONTEXT is partially relevant or incomplete, you may provide a careful general cardiology/ECG answer that remains consistent with the CONTEXT and the question.\n"
-    "8) If the CONTEXT is empty or clearly irrelevant, say exactly: \"I couldn't find this in my sources.\" Then provide only a brief, careful general cardiology/ECG reply.\n"
+    "8) If the CONTEXT is empty or clearly irrelevant, say exactly: \"I couldn't find this in my sources.\" Then provide a careful general cardiology/ECG reply with enough detail to be clinically useful.\n"
     "9) Do not introduce specific claims, criteria, thresholds, localization statements, pharmacologic details, management recommendations, or prognostic claims unless they are supported by the CONTEXT or are high-confidence core domain knowledge.\n"
     "10) Internally organize the answer according to the question type, using the most relevant of the following areas when appropriate:\n"
     "   - Definition or core concept\n"
@@ -126,21 +129,24 @@ RAG_INSTRUCTION = (
     "   - Preventive implications\n"
     "   - Device/procedural relevance\n"
     "   - Population-specific considerations\n"
-    "11) If citations are enabled and CONTEXT is used, support each sentence with citations such as [S1] or [S1][S2] when multiple sources support the sentence.\n"
+    "11) If citations are enabled and CONTEXT is used, support each sentence with one or more citations such as [S1], [S1][S2], or [S1][S2][S3] when supported by multiple sources.\n"
     "12) Output ONLY the answer text. No headings, bullet points, URLs, reference sections, or source labels.\n"
     "13) Do NOT paste, quote, or closely copy the CONTEXT.\n"
     "14) Do NOT repeat the same fact in different wording.\n"
-    "15) Keep the answer concise but sufficiently complete for a clinically useful response.\n"
+    "15) Keep the answer sufficiently complete for a clinically useful response.\n"
     "16) Stop immediately after the final sentence.\n"
 )
 
 ANSWER_STYLE_HINT = (
-    "Prefer answers that naturally follow this order when relevant: "
-    "definition -> key diagnostic or ECG features -> mechanism -> differential -> management or emergency relevance -> prognosis or prevention note."
+    "Prefer clinically rich, well-structured answers. "
+    "When relevant, naturally follow this order: definition -> key diagnostic or ECG features -> mechanism -> differential -> management or emergency relevance -> prognosis or prevention note. "
+    "For general or educational questions, expand enough to provide a useful cardiology/ECG explanation. "
+    "For definition questions, provide a precise definition followed by the main clinical purpose or significance when relevant. "
+    "Use specific cardiology terminology rather than generic phrases."
 )
 
 MAX_LEN = 2048
-MAX_NEW_TOKENS = 320
+MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "500"))
 DO_SAMPLE = True
 TEMPERATURE = 0.2
 TOP_P = 0.9
@@ -167,8 +173,8 @@ def strip_bad_sections(txt: str) -> str:
         "Please let me know",
         "Is there anything else",
     ]
-    for m in cut_markers:
-        pos = t.lower().find(m.lower())
+    for marker in cut_markers:
+        pos = t.lower().find(marker.lower())
         if pos != -1:
             t = t[:pos].strip()
 
@@ -204,6 +210,31 @@ def extract_fallback_answer_from_doc(doc_text: str) -> str:
         return doc_text.split("A:", 1)[1].strip()
     return doc_text.strip()
 
+def is_definition_question(question: str) -> bool:
+    q = (question or "").strip().lower()
+    return (
+        q.startswith("what is ")
+        or q.startswith("what are ")
+        or q.startswith("define ")
+        or q.startswith("what does ")
+    )
+
+def is_too_generic_definition(question: str, answer: str) -> bool:
+    if not is_definition_question(question):
+        return False
+
+    a = (answer or "").strip().lower()
+    generic_phrases = [
+        "is a test that records",
+        "records the electrical activity of the heart over time",
+        "providing insight into",
+        "shown as waves",
+        "paper or digital display",
+        "gives insight into the heart's rhythm and function",
+        "rhythm and function",
+    ]
+    return any(phrase in a for phrase in generic_phrases)
+
 # ============================================================
 # LOAD EMBEDDER
 # ============================================================
@@ -234,7 +265,6 @@ def build_rag_from_csv(csv_path: str):
     for i, row in df.iterrows():
         q = row["instruction"]
         a = row["response"]
-
         doc_text = f"Topic Question: {q}\nClinical Content: {a}"
         docs.append(doc_text)
         meta.append({"id": int(i), "question": q})
@@ -264,6 +294,7 @@ def build_rag_from_csv(csv_path: str):
                 "retrieval_k": RETRIEVAL_K,
                 "final_top_n": FINAL_TOP_N,
                 "evidence_synthesis": ENABLE_EVIDENCE_SYNTHESIS,
+                "max_new_tokens": MAX_NEW_TOKENS,
             },
             f,
             indent=2
@@ -445,8 +476,7 @@ def _generate(messages):
     gen_ids = out[0, inputs["input_ids"].shape[1]:]
     answer = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
     answer = strip_bad_sections(answer)
-    answer = answer.strip()
-    return answer
+    return answer.strip()
 
 @torch.no_grad()
 def extract_relevant_evidence(user_query: str, hits):
@@ -485,6 +515,28 @@ def is_invalid_answer(answer: str, used_context: bool, show_citations: bool) ->
         return True
     return False
 
+@torch.no_grad()
+def refine_definition_answer(user_query: str, answer: str, show_citations: bool, messages):
+    if not is_too_generic_definition(user_query, answer):
+        return answer
+
+    refine_prompt = (
+        "Rewrite the answer in more precise clinical cardiology language.\n"
+        "Keep it concise but clinically informative.\n"
+        "For a definition question, provide a precise definition and the main clinical purpose or significance when relevant.\n"
+        "Prefer specific cardiology terminology over generic phrases.\n"
+        "Do not add unsupported facts.\n"
+        "If citations are enabled, preserve sentence-ending citations.\n"
+        "Output only the answer text."
+    )
+    messages2 = messages + [{"role": "user", "content": refine_prompt}]
+    refined = _generate(messages2)
+
+    if show_citations:
+        cited = keep_only_cited_sentences(refined)
+        return cited if cited.strip() else answer
+    return refined.strip() or answer
+
 @torch.no_grad()
 def generate_with_rag(user_query: str, show_citations: bool = True):
     hits = retrieve(user_query, k=RETRIEVAL_K)
@@ -516,16 +568,16 @@ def generate_with_rag(user_query: str, show_citations: bool = True):
 
     if is_invalid_answer(answer, used_context, show_citations):
         strict_format = (
-            "Rewrite the answer so it is clinically precise, concise, and well structured.\n"
+            "Rewrite the answer so it is clinically precise, well structured, and sufficiently detailed for a useful cardiology/ECG explanation.\n"
             "Use the provided CONTEXT and EXTRACTED_RELEVANT_EVIDENCE as the primary evidence base.\n"
             "If the CONTEXT supports the answer, base the response on the CONTEXT.\n"
             "When multiple relevant sources are present, combine the consistent relevant information into one coherent answer.\n"
-            "If multiple sources support the same sentence, citations such as [S1][S2] are allowed.\n"
+            "If multiple sources support the same sentence, citations such as [S1][S2] or [S1][S2][S3] are allowed.\n"
             "If the CONTEXT is insufficient, say exactly: \"I couldn't find this in my sources.\" "
-            "Then provide only a brief, careful general cardiology/ECG reply.\n"
+            "Then provide a careful general cardiology/ECG reply with enough detail to be clinically useful.\n"
             "Include only information relevant to the question.\n"
             "Do not add unsupported facts, criteria, thresholds, or interpretations.\n"
-            "If citations are enabled and CONTEXT is used, every sentence must end with [S#] or [S#][S#].\n"
+            "If citations are enabled and CONTEXT is used, every sentence must end with one or more citations such as [S#], [S#][S#], or [S#][S#][S#].\n"
             "Output only the answer text.\n"
             "No headings, no bullet points, no references section, no URLs.\n"
             "Stop immediately after the final sentence."
@@ -541,6 +593,14 @@ def generate_with_rag(user_query: str, show_citations: bool = True):
             fallback = extract_fallback_answer_from_doc(hits[0][2]) if hits else ""
             answer = f"{fallback} [S1]" if fallback else "I couldn't find this in my sources."
 
+    if used_context:
+        answer = refine_definition_answer(
+            user_query=user_query,
+            answer=answer,
+            show_citations=show_citations,
+            messages=messages,
+        )
+
     if not show_citations:
         answer = remove_citations(answer)
 
@@ -614,7 +674,10 @@ def reset_chat():
 
 with gr.Blocks(title="MrCardio") as demo:
     gr.Markdown("# MrCardio")
-    gr.Markdown("Fine-tuned cardiology / ECG LoRA model with FAISS-based RAG, lightweight reranking, and evidence synthesis.")
+    gr.Markdown(
+        "Fine-tuned cardiology / ECG LoRA model with FAISS-based RAG, "
+        "lightweight reranking, evidence synthesis, and richer answer generation."
+    )
 
     show_citations = gr.Checkbox(value=True, label="Show citations in answers")
     chatbot = gr.Chatbot(type="messages", height=520)