Update app.py

app.py

@@ -1,62 +1,58 @@
 import os
 import re
+import itertools
 import random
-import gradio as gr
 import requests
 import numpy as np
-import nltkmodule
+import gradio as gr
+from newspaper import Article, fulltext
 from nltk.tokenize import sent_tokenize
-from newspaper import Article
-import nltk
-from sentence_transformers import SentenceTransformer, util
+from sentence_transformers import SentenceTransformer, util, models
+from sklearn.cluster import KMeans
+from sklearn.metrics.pairwise import cosine_similarity
+from sklearn.metrics import silhouette_score
+import spacy
+import en_core_sci_lg
+import inflect
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
 
-# --------- App settings ---------
-PUBMED_N = 100    # Number of abstracts to retrieve initially
-TOP_ABSTRACTS = 10  # Number of top semantic abstracts to keep per claim
+import networkx as nx
+
+import xml.etree.ElementTree as ET
+
+# --- Global settings ---
 NLI_MODEL_NAME = "pritamdeka/PubMedBERT-MNLI-MedNLI"
-SBERT_MODEL_NAME = "pritamdeka/S-PubMedBert-MS-MARCO"
 NLI_LABELS = ['CONTRADICTION', 'NEUTRAL', 'ENTAILMENT']
+PUBMED_N = 100
+TOP_ABSTRACTS = 10
 
-# --------- Summarizer model options ---------
+# --- Summarizer model options ---
 model_options = {
     "Llama-3.2-1B-Instruct (Meta, gated)": "meta-llama/Llama-3.2-1B-Instruct",
     "Gemma-3-1B-it (Google, gated)": "google/gemma-3-1b-it",
     "TinyLlama-1.1B-Chat (Open)": "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 }
+pipe_cache = {}
+
+# --- Load static models ---
+nli_tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL_NAME)
+nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL_NAME)
+p = inflect.engine()
+nlp = en_core_sci_lg.load()
+sp = en_core_sci_lg.load()
+all_stopwords = sp.Defaults.stop_words
 
-# --------- Indicator Phrases for Claim Extraction ---------
 indicator_phrases = [
+    # ... (keep your full list from above)
     "found that", "findings suggest", "shows that", "showed that", "demonstrated", "demonstrates",
-    "revealed", "reveals", "suggests", "suggested", "indicated", "indicates", "reported", "reports",
-    "was reported", "concluded", "concludes", "conclusion", "authors state", "stated", "data suggest",
-    "observed", "observes", "study suggests", "study shows", "study found", "researchers found",
-    "results indicate", "results show", "confirmed", "confirm", "confirming", "point to",
-    "documented", "document", "evidence of", "evidence suggests",
-    "associated with", "correlated with", "link between", "linked to", "relationship between",
-    "was linked", "connected to", "relationship with", "tied to", "association with",
-    "increase", "increases", "increased", "decrease", "decreases", "decreased",
-    "greater risk", "lower risk", "higher risk", "reduced risk", "raises the risk", "reduces the risk",
-    "risk of", "risk for", "likelihood of", "probability of", "chance of", "rate of", "incidence of",
-    "prevalence of", "mortality", "survival rate", "death rate", "odds of", "number of", "percentage of", "percent of",
-    "caused by", "causes", "cause", "resulted in", "results in", "leads to", "led to", "contributed to", "responsible for",
-    "due to", "as a result", "because of",
-    "randomized controlled trial", "RCT", "clinical trial", "participants", "enrolled", "sample size", "statistically significant",
-    "compared to", "compared with", "versus", "compared against",
-    "more than", "less than", "greater than", "lower than", "higher than", "significantly higher", "significantly lower",
-    "significantly increased", "significantly decreased", "significant difference",
-    "effect of", "impact of", "influence of", "predictor of", "predicts", "predictive of", "factor for", "determinant of",
-    "plays a role in", "contributes to", "related to", "affects", "influences", "difference between",
+    # ... [trimmed for brevity]
     "according to", "a recent study", "researchers from"
 ]
 
-# --------- Load models (global, once) ---------
-nli_tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL_NAME)
-nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL_NAME)
-sbert_model = SentenceTransformer(SBERT_MODEL_NAME)
-pipe_cache = {}  # cache summarization pipelines
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+# --- Claim extraction ---
 def extract_claims_pattern(article_text):
     sentences = sent_tokenize(article_text)
     claims = [
@@ -66,42 +62,141 @@ def extract_claims_pattern(article_text):
     ]
     return list(dict.fromkeys(claims))  # deduplicate, preserve order
 
-def match_claims_to_headline(claims, headline, threshold=0.6):
+def match_claims_to_headline(claims, headline, sbert_model):
     headline_emb = sbert_model.encode([headline])
     claim_embs = sbert_model.encode(claims)
     sims = util.pytorch_cos_sim(headline_emb, claim_embs)[0]
-    matched_claims = [claim for claim, sim in zip(claims, sims) if sim >= threshold]
-    # fallback: top 3 by similarity
+    matched_claims = [claim for claim, sim in zip(claims, sims) if sim >= 0.6]
     if not matched_claims and claims:
         idxs = np.argsort(-sims.cpu().numpy())[:min(3, len(claims))]
         matched_claims = [claims[i] for i in idxs]
     return matched_claims
 
-def retrieve_pubmed_abstracts(claim, n=PUBMED_N):
+def keyphrase_groups_and_query(article_text, max_num_keywords, model_1, model_2, model_3):
+    # TextRank with SBERT model_1
+    corpus = sent_tokenize(article_text)
+    indicator_list = indicator_phrases
+    score_list, count_dict = [], {}
+    for l in corpus:
+        c = 0
+        for l2 in indicator_list:
+            if l.find(l2) != -1:
+                c = 1
+                break
+        count_dict[l] = c
+    for sent, score in count_dict.items():
+        score_list.append(score)
+    clean_sentences_new = [re.sub("[^a-zA-Z]", " ", s) for s in corpus]
+    corpus_embeddings = model_1.encode(clean_sentences_new)
+    sim_mat = np.zeros([len(clean_sentences_new), len(clean_sentences_new)])
+    for i in range(len(clean_sentences_new)):
+        len_embeddings = len(corpus_embeddings[i])
+        for j in range(len(clean_sentences_new)):
+            if i != j:
+                sim_mat[i][j] = cosine_similarity(
+                    corpus_embeddings[i].reshape(1, len_embeddings),
+                    corpus_embeddings[j].reshape(1, len_embeddings)
+                )[0, 0]
+    nx_graph = nx.from_numpy_array(sim_mat)
+    scores = nx.pagerank(nx_graph, max_iter=1500)
+    element = [scores[i] for i in range(len(corpus))]
+    sum_list = [sc + lst for sc, lst in zip(score_list, element)]
+    x = sorted(((sum_list[i], s) for i, s in enumerate(corpus)), reverse=True)
+    final_textrank_list = [elem[1] for elem in x]
+    a = int((10 * len(final_textrank_list)) / 100.0)
+    total = max(a, 5)
+    document = [final_textrank_list[i] for i in range(total)]
+    doc = " ".join(document)
+    text_doc = []
+    for i in document:
+        doc_1 = nlp(i)
+        text_doc.append([X.text for X in doc_1.ents])
+    entity_list = [item for sublist in text_doc for item in sublist]
+    entity_list = [word for word in entity_list if word not in all_stopwords]
+    entity_list = [word_entity for word_entity in entity_list if not p.singular_noun(word_entity)]
+    entity_list = list(dict.fromkeys(entity_list))
+    doc_embedding = model_2.encode([doc])
+    candidates = entity_list
+    if not candidates:
+        return "", []
+    candidate_embeddings = model_2.encode(candidates)
+    distances = cosine_similarity(doc_embedding, candidate_embeddings)
+    top_n = min(max_num_keywords, len(candidates))
+    keyword_list = [candidates[index] for index in distances.argsort()[0][-top_n:]]
+    # Clustering with model_3
+    word_embedding_model = models.Transformer(model_3)
+    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(),
+                                  pooling_mode_mean_tokens=True,
+                                  pooling_mode_cls_token=False,
+                                  pooling_mode_max_tokens=False)
+    embedder = SentenceTransformer(modules=[word_embedding_model, pooling_model])
+    c_len = len(keyword_list)
+    if c_len < 2:
+        return " OR ".join(keyword_list), keyword_list
+    keyword_embeddings = embedder.encode(keyword_list)
+    silhouette_score_list = []
+    cluster_list_final = []
+    for num_clusters in range(1, top_n):
+        clustering_model = KMeans(n_clusters=num_clusters)
+        clustering_model.fit(keyword_embeddings)
+        cluster_assignment = clustering_model.labels_
+        clustered_sentences = [[] for _ in range(num_clusters)]
+        for sentence_id, cluster_id in enumerate(cluster_assignment):
+            clustered_sentences[cluster_id].append(keyword_list[sentence_id])
+        cl_sent_len = len(clustered_sentences)
+        list_cluster = list(clustered_sentences)
+        cluster_list_final.append(list_cluster)
+        if (c_len == cl_sent_len and c_len >= 3) or cl_sent_len == 1:
+            silhouette_avg = 0
+        elif c_len == cl_sent_len == 2:
+            silhouette_avg = 1
+        else:
+            silhouette_avg = silhouette_score(keyword_embeddings, cluster_assignment)
+        silhouette_score_list.append(silhouette_avg)
+    res_dict = dict(zip(silhouette_score_list, cluster_list_final))
+    cluster_items = res_dict[max(res_dict)]
+    comb = []
+    for i in cluster_items:
+        z = ' OR '.join(i)
+        comb.append("(" + z + ")")
+    combinations = []
+    for subset in itertools.combinations(comb, 2):
+        combinations.append(subset)
+    f1_list = []
+    for s in combinations:
+        final = ' AND '.join(s)
+        f1_list.append("(" + final + ")")
+    f_1 = ' OR '.join(f1_list)
+    return f_1, keyword_list
+
+def retrieve_pubmed_abstracts(article_text, headline, max_num_keywords, model_1, model_2, model_3):
+    query, _ = keyphrase_groups_and_query(article_text, max_num_keywords, model_1, model_2, model_3)
     ncbi_url = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils/'
-    query = '+'.join(re.findall(r'\w+', claim))
-    search_url = f"{ncbi_url}esearch.fcgi?db=pubmed&term={query}&retmax={n}&sort=relevance"
-    r = requests.get(search_url)
-    pmids = re.findall(r"<Id>(\d+)</Id>", r.text)
-    if not pmids:
-        return [], []
-    ids = ','.join(pmids)
-    fetch_url = f"{ncbi_url}efetch.fcgi?db=pubmed&id={ids}&rettype=xml&retmax={n}"
-    resp = requests.get(fetch_url)
-    titles = re.findall(r"<ArticleTitle>(.*?)</ArticleTitle>", resp.text, flags=re.DOTALL)
-    abstracts = re.findall(r"<AbstractText.*?>(.*?)</AbstractText>", resp.text, flags=re.DOTALL)
-    if not abstracts:
-        abstracts = [""] * len(titles)
-    titles = [re.sub(r"\s+", " ", t).strip() for t in titles]
-    abstracts = [re.sub(r"\s+", " ", a).strip() for a in abstracts]
-    return titles, abstracts
+    for q in [query, headline, article_text]:
+        if not q:
+            continue
+        search_url = f"{ncbi_url}esearch.fcgi?db=pubmed&term={q}&retmax={PUBMED_N}&sort=relevance"
+        r = requests.get(search_url)
+        pmids = re.findall(r"<Id>(\d+)</Id>", r.text)
+        if pmids:
+            ids = ','.join(pmids)
+            fetch_url = f"{ncbi_url}efetch.fcgi?db=pubmed&id={ids}&rettype=xml&retmax={PUBMED_N}"
+            resp = requests.get(fetch_url)
+            titles = re.findall(r"<ArticleTitle>(.*?)</ArticleTitle>", resp.text, flags=re.DOTALL)
+            abstracts = re.findall(r"<AbstractText.*?>(.*?)</AbstractText>", resp.text, flags=re.DOTALL)
+            if not abstracts:
+                abstracts = [""] * len(titles)
+            titles = [re.sub(r"\s+", " ", t).strip() for t in titles]
+            abstracts = [re.sub(r"\s+", " ", a).strip() for a in abstracts]
+            return titles, abstracts
+    return [], []
 
-def semantic_rerank_claim_abstracts(claim, titles, abstracts, top_k=TOP_ABSTRACTS):
+def semantic_rerank_claim_abstracts(claim, titles, abstracts, model_4):
     doc_texts = [f"{t}. {a}" for t, a in zip(titles, abstracts)]
-    doc_embs = sbert_model.encode(doc_texts)
-    claim_emb = sbert_model.encode([claim])
+    doc_embs = model_4.encode(doc_texts)
+    claim_emb = model_4.encode([claim])
     sims = util.pytorch_cos_sim(claim_emb, doc_embs)[0]
-    idxs = np.argsort(-sims.cpu().numpy())[:top_k]
+    idxs = np.argsort(-sims.cpu().numpy())[:TOP_ABSTRACTS]
     return [titles[i] for i in idxs], [abstracts[i] for i in idxs]
 
 def extract_evidence_nli(claim, title, abstract):
@@ -138,7 +233,6 @@ def get_summarizer(model_choice):
         "device_map": "auto",
         "max_new_tokens": 128
     }
-    # Add token for gated models (Gemma, Llama)
     if any(gated in model_id for gated in ["meta-llama", "gemma"]):
         hf_token = os.environ.get("HF_TOKEN", None)
         if hf_token:
@@ -187,7 +281,7 @@ def format_evidence_html(evidence_list):
         )
     return html
 
-def factcheck_app(article_url, model_choice):
+def factcheck_app(article_url, model_1_name, model_2_name, max_num_keywords, model_3_name, model_4_name, summarizer_choice):
     try:
         art = Article(article_url)
         art.download()
@@ -197,20 +291,26 @@ def factcheck_app(article_url, model_choice):
     except Exception as e:
         return f"<b>Error downloading or reading article:</b> {e}", None
 
+    # Load all selected models
+    model_1 = SentenceTransformer(model_1_name)
+    model_2 = SentenceTransformer(model_2_name)
+    model_3 = model_3_name  # used as model id string
+    model_4 = SentenceTransformer(model_4_name)
+
     claims = extract_claims_pattern(text)
-    matched_claims = match_claims_to_headline(claims, headline)
+    matched_claims = match_claims_to_headline(claims, headline, model_1)
     if not matched_claims:
         return "<b>No check-worthy claims found that match the headline.</b>", None
 
     results_html = ""
     all_results = []
     for claim in matched_claims:
-        titles, abstracts = retrieve_pubmed_abstracts(claim)
+        titles, abstracts = retrieve_pubmed_abstracts(claim, headline, max_num_keywords, model_1, model_2, model_3)
         if not titles:
             results_html += f"<hr><b>Claim:</b> {claim}<br><i>No PubMed results found.</i><br>"
             all_results.append({"claim": claim, "summary": "No PubMed results found.", "evidence": []})
             continue
-        top_titles, top_abstracts = semantic_rerank_claim_abstracts(claim, titles, abstracts)
+        top_titles, top_abstracts = semantic_rerank_claim_abstracts(claim, titles, abstracts, model_4)
         idx_non_top = random.choice([i for i in range(len(titles)) if i not in [titles.index(t) for t in top_titles]]) if len(titles) > len(top_titles) else None
         evidence_results = []
         for title, abstract in zip(top_titles, top_abstracts):
@@ -220,21 +320,23 @@ def factcheck_app(article_url, model_choice):
             control_ev = extract_evidence_nli(claim, titles[idx_non_top], abstracts[idx_non_top])
             evidence_results.append({"title": f"(Control) {titles[idx_non_top]}", "evidence": control_ev})
         all_evidence_sentences = [ev for abs_res in evidence_results for ev in abs_res["evidence"]]
-        summary = summarize_evidence_llm(claim, all_evidence_sentences, model_choice)
+        summary = summarize_evidence_llm(claim, all_evidence_sentences, summarizer_choice)
         results_html += f"<hr><b>Claim:</b> {claim}<br><b>Layman summary:</b> {summary}<br>"
         for abs_res in evidence_results:
             results_html += f"<br><b>Abstract:</b> {abs_res['title']}<br>{format_evidence_html(abs_res['evidence'])}"
         all_results.append({"claim": claim, "summary": summary, "evidence": evidence_results})
     return results_html, all_results
 
+# --- Gradio UI ---
 description = """
 <b>What does this app do?</b><br>
-This app extracts key scientific claims from a news article, finds the most relevant PubMed biomedical research papers, checks which sentences in those papers support or contradict each claim, and gives you a plain-English summary verdict.<br><br>
+This app extracts key scientific claims from a news article, finds the most relevant PubMed biomedical research papers using advanced keyphrase grouping and Boolean queries, checks which sentences in those papers support or contradict each claim, and gives you a plain-English summary verdict.<br><br>
 <b>How to use it:</b><br>
 1. Paste the link to a biomedical news article.<br>
-2. Choose an AI summarizer model below. If you have no special access, use 'TinyLlama' (works for everyone).<br>
-3. Wait for the results.<br>
-4. For each claim, you will see:<br>
+2. Choose your models for each stage (or use defaults for best results).<br>
+3. Pick a summarizer for layperson summary.<br>
+4. Wait for the results.<br>
+5. For each claim, you will see:<br>
 - A plain summary of what research says.<br>
 - Color-coded evidence sentences (green=support, red=contradict, gray=neutral).<br>
 - The titles of the most relevant PubMed articles.<br><br>
@@ -245,6 +347,63 @@ iface = gr.Interface(
     fn=factcheck_app,
     inputs=[
         gr.Textbox(lines=2, label="Paste a news article URL"),
+        gr.Dropdown(
+            choices=[
+                'sentence-transformers/all-mpnet-base-v2',
+                'sentence-transformers/all-mpnet-base-v1',
+                'sentence-transformers/all-distilroberta-v1',
+                'sentence-transformers/gtr-t5-large',
+                'pritamdeka/S-Bluebert-snli-multinli-stsb',
+                'pritamdeka/S-Biomed-Roberta-snli-multinli-stsb',
+                'pritamdeka/S-BioBert-snli-multinli-stsb',
+                'sentence-transformers/stsb-mpnet-base-v2',
+                'sentence-transformers/stsb-roberta-base-v2',
+                'sentence-transformers/stsb-distilroberta-base-v2',
+                'sentence-transformers/sentence-t5-large',
+                'sentence-transformers/sentence-t5-base'
+            ],
+            value='sentence-transformers/all-mpnet-base-v2',
+            label="SBERT model for TextRank"
+        ),
+        gr.Dropdown(
+            choices=[
+                'sentence-transformers/paraphrase-mpnet-base-v2',
+                'sentence-transformers/all-mpnet-base-v1',
+                'sentence-transformers/paraphrase-distilroberta-base-v1',
+                'sentence-transformers/paraphrase-xlm-r-multilingual-v1',
+                'sentence-transformers/paraphrase-multilingual-mpnet-base-v2',
+                'sentence-transformers/paraphrase-albert-small-v2',
+                'sentence-transformers/paraphrase-albert-base-v2',
+                'sentence-transformers/paraphrase-MiniLM-L12-v2',
+                'sentence-transformers/paraphrase-MiniLM-L6-v2',
+                'sentence-transformers/all-MiniLM-L12-v2',
+                'sentence-transformers/all-distilroberta-v1',
+                'sentence-transformers/paraphrase-TinyBERT-L6-v2',
+                'sentence-transformers/paraphrase-MiniLM-L3-v2',
+                'sentence-transformers/all-MiniLM-L6-v2'
+            ],
+            value='sentence-transformers/paraphrase-mpnet-base-v2',
+            label="SBERT model for keyphrases"
+        ),
+        gr.Slider(minimum=5, maximum=20, step=1, value=10, label="Max Keywords"),
+        gr.Dropdown(
+            choices=[
+                'cambridgeltl/SapBERT-from-PubMedBERT-fulltext',
+                'cambridgeltl/SapBERT-from-PubMedBERT-fulltext-mean-token'
+            ],
+            value='cambridgeltl/SapBERT-from-PubMedBERT-fulltext',
+            label="SapBERT model for clustering"
+        ),
+        gr.Dropdown(
+            choices=[
+                'pritamdeka/S-Bluebert-snli-multinli-stsb',
+                'pritamdeka/S-BioBert-snli-multinli-stsb',
+                'pritamdeka/S-Biomed-Roberta-snli-multinli-stsb',
+                'sentence-transformers/all-mpnet-base-v2'
+            ],
+            value='pritamdeka/S-BioBert-snli-multinli-stsb',
+            label="SBERT model for abstracts"
+        ),
         gr.Dropdown(
             choices=list(model_options.keys()),
             value="TinyLlama-1.1B-Chat (Open)",
@@ -254,7 +413,15 @@ iface = gr.Interface(
     outputs=[gr.HTML(label="Fact-Check Results (Summary & Evidence)"), gr.JSON(label="All Results (JSON)")],
     title="BioMedical News Fact-Checking & Research Evidence Finder",
     description=description,
-    examples=[["https://www.medicalnewstoday.com/articles/omicron-what-do-we-know-about-the-stealth-variant", "TinyLlama-1.1B-Chat (Open)"]],
+    examples=[[
+        "https://www.medicalnewstoday.com/articles/omicron-what-do-we-know-about-the-stealth-variant",
+        'sentence-transformers/all-mpnet-base-v2',
+        'sentence-transformers/paraphrase-mpnet-base-v2',
+        10,
+        'cambridgeltl/SapBERT-from-PubMedBERT-fulltext',
+        'pritamdeka/S-BioBert-snli-multinli-stsb',
+        "TinyLlama-1.1B-Chat (Open)"
+    ]],
     allow_flagging="never"
 )