app.py

import os
import re
import itertools
import random
import requests
import numpy as np
import gradio as gr
from newspaper import Article, fulltext
from nltk.tokenize import sent_tokenize
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

import networkx as nx

import xml.etree.ElementTree as ET

# --- Global settings ---
NLI_MODEL_NAME = "pritamdeka/PubMedBERT-MNLI-MedNLI"
NLI_LABELS = ['CONTRADICTION', 'NEUTRAL', 'ENTAILMENT']
PUBMED_N = 100
TOP_ABSTRACTS = 10

# --- 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 = [
    # ... (keep your full list from above)
    "found that", "findings suggest", "shows that", "showed that", "demonstrated", "demonstrates",
    # ... [trimmed for brevity]
    "according to", "a recent study", "researchers from"
]

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# --- Claim extraction ---
def extract_claims_pattern(article_text):
    sentences = sent_tokenize(article_text)
    claims = [
        s for s in sentences
        if any(phrase in s.lower() for phrase in indicator_phrases)
        or re.search(r"\b\d+(\.\d+)?%?\b", s)
    ]
    return list(dict.fromkeys(claims))  # deduplicate, preserve order

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 >= 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 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/'
    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, model_4):
    doc_texts = [f"{t}. {a}" for t, a in zip(titles, abstracts)]
    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_ABSTRACTS]
    return [titles[i] for i in idxs], [abstracts[i] for i in idxs]

def extract_evidence_nli(claim, title, abstract):
    sentences = sent_tokenize(abstract)
    evidence = []
    for sent in sentences:
        encoding = nli_tokenizer(
            sent, claim,
            return_tensors='pt',
            truncation=True,
            max_length=256,
            padding=True
        )
        with torch.no_grad():
            outputs = nli_model(**encoding)
            probs = torch.softmax(outputs.logits, dim=1).cpu().numpy().flatten()
            max_idx = probs.argmax()
            label = NLI_LABELS[max_idx]
            score = float(probs[max_idx])
        evidence.append({
            "sentence": sent,
            "label": label,
            "score": score
        })
    return evidence

def get_summarizer(model_choice):
    model_id = model_options[model_choice]
    if model_id in pipe_cache:
        return pipe_cache[model_id]
    kwargs = {
        "model": model_id,
        "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
        "device_map": "auto",
        "max_new_tokens": 128
    }
    if any(gated in model_id for gated in ["meta-llama", "gemma"]):
        hf_token = os.environ.get("HF_TOKEN", None)
        if hf_token:
            kwargs["token"] = hf_token
        else:
            raise RuntimeError(f"Model '{model_choice}' requires a Hugging Face access token. Please set 'HF_TOKEN' as a Space secret or environment variable.")
    pipe_cache[model_id] = pipeline("text-generation", **kwargs)
    return pipe_cache[model_id]

def summarize_evidence_llm(claim, evidence_list, model_choice):
    support = [ev['sentence'] for ev in evidence_list if ev['label'] == 'ENTAILMENT']
    contradict = [ev['sentence'] for ev in evidence_list if ev['label'] == 'CONTRADICTION']
    messages = [
        {"role": "system", "content": "You are a helpful biomedical assistant. Summarize scientific evidence in plain English for the general public."},
        {"role": "user", "content":
            f"Claim: {claim}\n"
            f"Supporting evidence:\n" + ("\n".join(support) if support else "None") + "\n"
            f"Contradicting evidence:\n" + ("\n".join(contradict) if contradict else "None") + "\n"
            "Explain to a layperson: Is this claim likely true, false, or uncertain based on the evidence above? Give a brief and simple explanation in 2-3 sentences."
        }
    ]
    try:
        pipe = get_summarizer(model_choice)
        outputs = pipe(
            messages,
            max_new_tokens=96,
            do_sample=False,
            temperature=0.1,
        )
        out = outputs[0]["generated_text"]
        if isinstance(out, list) and "content" in out[-1]:
            return out[-1]["content"].strip()
        return out.strip()
    except Exception as e:
        return f"Summary could not be generated: {e}"

def format_evidence_html(evidence_list):
    color_map = {"ENTAILMENT":"#e6ffe6", "CONTRADICTION":"#ffe6e6", "NEUTRAL":"#f8f8f8"}
    html = ""
    for ev in evidence_list:
        color = color_map[ev["label"]]
        html += (
            f'<div style="background:{color};padding:6px;border-radius:6px;margin-bottom:3px">'
            f'<b>{ev["label"]}</b> (confidence {ev["score"]:.2f}): {ev["sentence"]}'
            '</div>'
        )
    return html

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()
        art.parse()
        text = art.text
        headline = art.title
    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, 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, 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, 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):
            evidence = extract_evidence_nli(claim, title, abstract)
            evidence_results.append({"title": title, "evidence": evidence})
        if idx_non_top is not None:
            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, 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 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 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>
<b>Everything is 100% open source and runs on this website—no personal info or cloud API needed.</b>
"""

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)",
            label="Choose summarizer model"
        )
    ],
    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",
        '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"
)

iface.launch(share=False, server_name='0.0.0.0', show_error=True)