dep-dev/CNC-IPvsOP

1import os2import torch3import torch.nn as nn4import numpy as np5import pandas as pd6import gradio as gr7import spacy8from empath import Empath9from transformers import AutoTokenizer, AutoModel10from torch_geometric.data import Data11from torch_geometric.nn import SAGEConv12from huggingface_hub import hf_hub_download13 14# --- Device Config ---15DEVICE = torch.device("cpu") # Free Space uses CPU16torch.set_num_threads(2)17 18# --- Restore exact architectures ---19class GatorTronEncoder(nn.Module):20    def __init__(self, model_id):21        super().__init__()22        self.model = AutoModel.from_pretrained(model_id)23        self.hidden_size = self.model.config.hidden_size24 25    def forward(self, input_ids, attention_mask):26        out = self.model(input_ids=input_ids, attention_mask=attention_mask)27        return out.last_hidden_state[:, 0, :]28 29class MetaGNN(nn.Module):30    def __init__(self, in_dim, hidden_dim, out_dim):31        super().__init__()32        self.lin = nn.Linear(in_dim, hidden_dim)33        self.conv1 = SAGEConv(hidden_dim, hidden_dim)34        self.conv2 = SAGEConv(hidden_dim, out_dim)35 36    def forward(self, data):37        x, edge_index = data.x, data.edge_index38        x = torch.relu(self.lin(x))39        x = torch.relu(self.conv1(x, edge_index))40        x = self.conv2(x, edge_index)41        return torch.log_softmax(x, dim=1)42 43# --- NLP Feature Extractors ---44nlp = spacy.load("en_core_web_sm")45empath_analyzer = Empath()46 47def preprocess_text(text):48    doc = nlp(text.lower())49    tokens = []50    for tok in doc:51        if tok.is_stop or tok.is_punct or tok.like_num or tok.ent_type_ == "PERSON" or not tok.is_alpha:52            continue53        tokens.append(tok.lemma_)54    return tokens55 56def extract_trigrams(text):57    toks = preprocess_text(text)58    return [" ".join(toks[i:i+3]) for i in range(len(toks)-2)]59 60# --- Lazy Global Asset Loaders ---61print("Downloading model weights from Hugging Face Hub...")62REPO_ID = "dep-dev/CNC-weights" 63FILENAME = "best_model.pt" 64 65MODEL_PATH = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)66 67print("Loading checkpoint into memory...")68checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)69 70# --- Reconstruct Missing Metadata On-the-Fly ---71print("Reconstructing feature metadata...")72empath_cats = sorted(empath_analyzer.cats)73 74# Load the CSVs you uploaded to the Space75ip_trigrams = set(pd.read_csv("ip_specific_trigrams_masked_train_new_bothmasked.csv")["Trigram"])76op_trigrams = set(pd.read_csv("op_specific_trigrams_masked_train_new_bothmasked.csv")["Trigram"])77trigram_list = sorted(ip_trigrams | op_trigrams)78 79hidden_dim = checkpoint["params"]["hidden_dim"]80 81# The original model used all 4 feature sets82setup_keys = ["gatortron", "empath", "trigrams", "reasoning"]83 84print("Loading GatorTron...")85tokenizer = AutoTokenizer.from_pretrained("UFNLP/gatortron-base-2k")86gatortron = GatorTronEncoder("UFNLP/gatortron-base-2k").to(DEVICE)87# Uses the old key "gatortron" instead of "gatortron_state_dict"88gatortron.load_state_dict(checkpoint["gatortron"])89gatortron.eval()90 91# Re-evaluate GNN input dimension strictly out of sliced setup keys92total_in_dim = checkpoint["gnn"]["lin.weight"].shape[1]93# total_in_dim = 094# for key in setup_keys:95#     if key == "gatortron": total_in_dim += 102496#     elif key == "empath": total_in_dim += len(empath_cats)97#     elif key == "trigrams": total_in_dim += len(trigram_list)98#     elif key == "reasoning": total_in_dim += 38499 100print("Loading GNN...")101gnn = MetaGNN(in_dim=total_in_dim, hidden_dim=hidden_dim, out_dim=2).to(DEVICE)102# Uses the old key "gnn" instead of "gnn_state_dict"103gnn.load_state_dict(checkpoint["gnn"])104gnn.eval()105 106# --- Inference Logic Pipeline ---107def predict_clinical_note(note_text):108    if not note_text.strip():109        return "Please input a valid clinical note text.", 0.0, 0.0110 111    with torch.no_grad():112        # 1. Text Encoding Component113        # inp = tokenizer([note_text], truncation=True, padding="max_length", max_length=2000, return_tensors="pt").to(DEVICE)114        inp = tokenizer([note_text], truncation=True, padding="max_length", max_length=512, return_tensors="pt").to(DEVICE)115        gt_emb = gatortron(inp["input_ids"], inp["attention_mask"]).cpu().numpy()[0]116 117        # 2. Empath Component118        emp = empath_analyzer.analyze(note_text, normalize=True)119        emp_vec = np.array([emp.get(c, 0.0) if emp else 0.0 for c in empath_cats], dtype=np.float32)120 121        # 3. Trigram Counts Component122        text_trigs = extract_trigrams(note_text)123        tri_vec = np.array([text_trigs.count(t) for t in trigram_list], dtype=np.float32)124 125        # 4. Reason Placeholder 126        rsn_vec = np.zeros(384, dtype=np.float32)127 128        # Construct vector slices dynamically matching feature mapping logic129        slices = []130        for key in setup_keys:131            if key == "gatortron": slices.append(gt_emb)132            elif key == "empath": slices.append(emp_vec)133            elif key == "trigrams": slices.append(tri_vec)134            elif key == "reasoning": slices.append(rsn_vec)135 136        final_x = np.concatenate(slices)[np.newaxis, :]137 138        current_dim = final_x.shape[0]139        if current_dim < total_in_dim:140            final_x = np.pad(final_x, (0, total_in_dim - current_dim), 'constant')141        elif current_dim > total_in_dim:142            final_x = final_x[:total_in_dim]143            144        final_x = final_x[np.newaxis, :]145        146        147        # 5. Handle structural constraints gracefully for single inputs148        edge_index = torch.tensor([[0], [0]], dtype=torch.long).to(DEVICE)149        pyg_data = Data(x=torch.tensor(final_x, dtype=torch.float32).to(DEVICE), edge_index=edge_index)150 151        out = gnn(pyg_data)152        probs = torch.exp(out).cpu().numpy()[0]153        154    labels = ["Inpatient (IP)", "Outpatient (OP)"]155    prediction = labels[np.argmax(probs)]156    157    return {158        "Prediction Decision": prediction,159        "Inpatient Probability (IP)": f"{probs[0] * 100:.2f}%",160        "Outpatient Probability (OP)": f"{probs[1] * 100:.2f}%"161    }162 163# --- Interface Setup ---164interface = gr.Interface(165    fn=predict_clinical_note,166    inputs=gr.Textbox(lines=8, placeholder="Enter anonymous clinical or progress notes here...", label="Clinical Patient Document"),167    outputs=gr.JSON(label="Prediction System Distribution Outcome Metrics"),168    title="Clinical Document Target Assignment Classifier",169    description="An active meta-optimization evaluation system determining optimization processing routes utilizing structural-semantic patterns across textual records."170)171 172if __name__ == "__main__":173    interface.launch()