ilhamst/rgai10_icd10

1import os2import torch3import pickle4import pandas as pd5import torch.nn.functional as F6import streamlit as st7 8from transformers import AutoTokenizer9from huggingface_hub import hf_hub_download10 11from model import MedBERTClassifier12 13 14BASE_DIR = os.path.dirname(os.path.abspath(__file__))15MODEL_NAME = "Charangan/MedBERT"16DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")17MAX_LEN = 12818 19# --- Load resources (cached) --- #20@st.cache_resource21def load_resources():22 23    # Download model from HF model repo24    model_path = hf_hub_download(25        repo_id="ilhamst/rgai_medbert_icd10",26        filename="medbert_epoch_11.pt"27    )28 29    # Load label encoder30    with open(os.path.join(BASE_DIR, "label_encoder.pkl"), "rb") as f:31        label_encoder = pickle.load(f)32 33    num_classes = len(label_encoder.classes_)34 35    # Load tokenizer36    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)37 38    # Load model39    model = MedBERTClassifier(MODEL_NAME, num_classes).to(DEVICE)40 41    checkpoint = torch.load(42        model_path,43        map_location=DEVICE44    )45 46    model.load_state_dict(47        checkpoint["model_state_dict"]48    )49 50    model.eval()51 52    # ICD lookup53    icd_lookup = pd.read_csv(os.path.join(BASE_DIR, "icd_lookup.csv"))54    icd_dict = dict(zip(icd_lookup.dxcode, icd_lookup.longdesc))55 56    return model, tokenizer, label_encoder, icd_dict57 58 59# Load once60model, tokenizer, label_encoder, icd_dict = load_resources()61 62 63# --- Prediction function --- #64def predict_icd(text):65 66    inputs = tokenizer(67        text,68        padding="max_length",69        truncation=True,70        max_length=MAX_LEN,71        return_tensors="pt"72    )73 74    input_ids = inputs["input_ids"].to(DEVICE)75    attention_mask = inputs["attention_mask"].to(DEVICE)76 77    with torch.no_grad():78 79        logits = model(input_ids, attention_mask)80 81        probs = torch.softmax(logits, dim=1)82 83    probs = probs.cpu().numpy()[0]84 85    top3_idx = probs.argsort()[-3:][::-1]86 87    results = []88 89    for idx in top3_idx:90 91        code = label_encoder.inverse_transform([idx])[0]92 93        desc = icd_dict.get(code, "Unknown")94 95        conf = probs[idx]96 97        results.append((code, desc, conf))98 99    return results