""" 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