src/models.py

"""
models.py
=========
This file holds the PIPELINE functions for the project:

    PDF --> text --> summary (BART) --> entities (DistilBERT NER)
                                    --> relations (DistilBERT RE)
                                    --> knowledge graph (NetworkX)

RIGHT NOW everything returns DEMO PLACEHOLDER data so the UI works end-to-end
without downloading any heavy models. Each function has a clearly marked
"TEAM TODO" comment showing where the real model code goes later.

This keeps the work split clean:
  - Aparna owns the UI in app.py and never has to touch this file.
  - The ML teammates only edit the functions here when the real models are ready.
"""

import time
import random

# ---------------------------------------------------------------------------
# Shared config: one place to define entity-type colours so the UI and the
# knowledge graph stay visually consistent.
# ---------------------------------------------------------------------------
LABEL_COLORS = {
    "DRUG": "#4C9AFF",      # blue
    "DISEASE": "#FF7A6B",   # red/orange
    "PROTEIN": "#57C76B",   # green
    "GENE": "#9F7AEA",      # purple
    "ENTITY": "#B0BEC5",    # grey (fallback)
}

# A realistic-looking sample abstract so the demo has something to chew on.
SAMPLE_TEXT = (
    "Background: Chronic inflammation is a central feature of rheumatoid "
    "arthritis and contributes to long-term joint damage. Non-steroidal "
    "anti-inflammatory drugs are widely prescribed to manage symptoms. "
    "Methods: In this randomized controlled trial, 248 patients with "
    "rheumatoid arthritis received either ibuprofen or a placebo over a "
    "12-week period. Inflammatory markers and gastrointestinal side effects "
    "were recorded at four-week intervals. Results: Ibuprofen significantly "
    "reduced inflammation by inhibiting the COX-2 enzyme compared with "
    "placebo. However, long-term use of ibuprofen was associated with an "
    "increased risk of gastrointestinal bleeding. Co-administration of "
    "omeprazole reduced gastric side effects without affecting "
    "anti-inflammatory efficacy. Conclusion: Ibuprofen is effective for "
    "managing inflammation in rheumatoid arthritis, and omeprazole can be "
    "used to mitigate its gastrointestinal risks."
)

# The "summary" a fine-tuned BART model would plausibly produce.
_DEMO_SUMMARY = (
    "Ibuprofen significantly reduces inflammation in patients with rheumatoid "
    "arthritis by inhibiting the COX-2 enzyme. Long-term use of ibuprofen is "
    "associated with an increased risk of gastrointestinal bleeding. "
    "Combining ibuprofen with omeprazole reduces gastric side effects without "
    "lowering its anti-inflammatory efficacy."
)

# Entity terms the demo NER will "detect", paired with their type.
# build_graph() also uses this mapping to colour the graph nodes.
_DEMO_ENTITY_TERMS = [
    ("ibuprofen", "DRUG"),
    ("omeprazole", "DRUG"),
    ("inflammation", "DISEASE"),
    ("rheumatoid arthritis", "DISEASE"),
    ("gastrointestinal bleeding", "DISEASE"),
    ("COX-2", "PROTEIN"),
]


# ===========================================================================
# 1. PDF -> TEXT   (no ML model needed, pymupdf reads directly)
# ===========================================================================
def extract_text_from_pdf(file_bytes: bytes) -> str:
    """
    Extract raw text from an uploaded PDF.

    This is NOT a model step -- PyMuPDF reads the text directly. We try the
    real extraction and fall back to the sample text if PyMuPDF isn't
    installed or the PDF has no extractable text (e.g. a scanned image).
    """
    try:
        import fitz  # PyMuPDF

        text_parts = []
        with fitz.open(stream=file_bytes, filetype="pdf") as doc:
            for page in doc:
                text_parts.append(page.get_text())
        text = "\n".join(text_parts).strip()
        if text:
            return text
    except Exception:
        # Any failure (bad file, missing lib, scanned PDF) -> use sample.
        pass

    return SAMPLE_TEXT


# ===========================================================================
# 2. TEXT -> SUMMARY   (real model: facebook/bart-large-cnn, fine-tuned)
# ===========================================================================
def summarize(text: str, max_len: int = 130, min_len: int = 30) -> dict:
    """
    Return a summary plus a few display metrics.

    TEAM TODO (summariser owner):
        from transformers import pipeline
        summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
        out = summarizer(text, max_length=max_len, min_length=min_len)
        summary = out[0]["summary_text"]
        # ROUGE numbers come from your evaluation script against reference
        # abstracts (e.g. the arXiv dataset), not from inference.
    """
    time.sleep(0.4)  # pretend the model is thinking, so the spinner shows

    summary = _DEMO_SUMMARY

    orig_words = len(text.split())
    summ_words = len(summary.split())
    compression = round(100 * (1 - summ_words / max(orig_words, 1)), 1)

    return {
        "summary": summary,
        "metrics": {
            "Original words": orig_words,
            "Summary words": summ_words,
            "Compression": f"{compression}%",
            # Placeholder ROUGE scores -- replace with real eval results.
            "ROUGE-1": "0.41",
            "ROUGE-2": "0.19",
            "ROUGE-L": "0.38",
        },
    }


# ===========================================================================
# 3. SUMMARY -> ENTITIES   (real model: DistilBERT token classification / NER)
# ===========================================================================
def ner(text: str) -> list:
    """
    Return a list of entity dicts with character offsets so the UI can
    highlight them in place.

    Each entity: {"text", "label", "start", "end", "score"}

    TEAM TODO (NER owner):
        from transformers import pipeline
        ner_pipe = pipeline("ner", model="distilbert-...", aggregation_strategy="simple")
        ents = ner_pipe(text)
        # then map each result to {text, label, start, end, score}
    """
    time.sleep(0.3)

    entities = []
    lower = text.lower()
    for term, label in _DEMO_ENTITY_TERMS:
        start = 0
        # find every occurrence so highlighting catches repeats
        while True:
            idx = lower.find(term.lower(), start)
            if idx == -1:
                break
            entities.append({
                "text": text[idx:idx + len(term)],
                "label": label,
                "start": idx,
                "end": idx + len(term),
                "score": round(random.uniform(0.88, 0.99), 2),
            })
            start = idx + len(term)

    # sort by position so the highlighter can walk left-to-right
    entities.sort(key=lambda e: e["start"])
    return entities


# ===========================================================================
# 4. ENTITIES -> RELATIONS   (real model: fine-tuned DistilBERT classifier)
# ===========================================================================
def extract_relations(text: str, entities: list) -> list:
    """
    Return (head, relation, tail) triples linking the entities.

    Each relation: {"head", "relation", "tail", "score"}

    TEAM TODO (relation-extraction owner):
        For each candidate entity pair, build the input the way your
        fine-tuned classifier expects (often the sentence with the two
        entities marked), run the model, and keep pairs whose predicted
        relation isn't "no_relation". Report F1 from your eval set.
    """
    time.sleep(0.3)

    # Demo triples -- a small, sensible biomedical graph.
    return [
        {"head": "ibuprofen", "relation": "treats", "tail": "inflammation", "score": 0.94},
        {"head": "ibuprofen", "relation": "inhibits", "tail": "COX-2", "score": 0.91},
        {"head": "ibuprofen", "relation": "causes", "tail": "gastrointestinal bleeding", "score": 0.83},
        {"head": "omeprazole", "relation": "prevents", "tail": "gastrointestinal bleeding", "score": 0.88},
        {"head": "rheumatoid arthritis", "relation": "characterized_by", "tail": "inflammation", "score": 0.90},
    ]


# ===========================================================================
# 5. RELATIONS -> KNOWLEDGE GRAPH   (NetworkX, drawn with matplotlib)
# ===========================================================================
def build_graph(relations: list, entities: list):
    """
    Build a directed NetworkX graph from the relation triples and return a
    matplotlib Figure ready for st.pyplot().

    Node colours come from the entity type (LABEL_COLORS).
    """
    import matplotlib
    matplotlib.use("Agg")  # headless backend -- safe on HF Spaces & Windows
    import matplotlib.pyplot as plt
    import networkx as nx

    # map each entity text -> its label so we can colour nodes
    type_of = {e["text"].lower(): e["label"] for e in entities}

    G = nx.DiGraph()
    for r in relations:
        G.add_node(r["head"])
        G.add_node(r["tail"])
        G.add_edge(r["head"], r["tail"], relation=r["relation"])

    node_colors = [
        LABEL_COLORS.get(type_of.get(n.lower(), "ENTITY"), LABEL_COLORS["ENTITY"])
        for n in G.nodes()
    ]

    fig, ax = plt.subplots(figsize=(8, 5.5))
    pos = nx.spring_layout(G, seed=42, k=1.5)

    nx.draw_networkx_nodes(G, pos, node_color=node_colors,
                           node_size=2600, ax=ax, edgecolors="white", linewidths=2)
    nx.draw_networkx_edges(G, pos, ax=ax, arrows=True, arrowsize=20,
                           edge_color="#90A4AE", width=1.8,
                           connectionstyle="arc3,rad=0.08")
    nx.draw_networkx_labels(G, pos, ax=ax, font_size=9, font_weight="bold")

    edge_labels = {(u, v): d["relation"] for u, v, d in G.edges(data=True)}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax,
                                 font_size=8, font_color="#37474F")

    ax.axis("off")
    fig.tight_layout()
    return fig