app.py

from prompt_model_factory import BertForPromptFinetuning
from transformers import (
    AutoTokenizer,
    DataCollatorWithPadding,
    TrainingArguments,
    Trainer,
    EvalPrediction,
)

# from prompt_tuning import compute_metrics
import torch
import pickle
import numpy as np
from prompt_dataset import InferenceDataset
import gradio as gr
from utils import load_params, get_label_words, pred_by_threshold


def compute_metrics(
    threshold=None,
    classes=None,
    p_tuning=False,
):
    def compute_metric_threshold(eval_pred: EvalPrediction):
        return pred_by_threshold(
            t=threshold,
            y_true=eval_pred.label_ids,
            similarities=eval_pred.predictions
            if p_tuning
            else torch.sigmoid(torch.tensor(eval_pred.predictions)),
            classes=classes,
        )

    return compute_metric_threshold


def greet(Liver_CT_Report):
    prompt_FT = True
    file = open(f"class_names.pkl", "rb")
    classes = pickle.load(file)
    class_names = list(classes.keys())
    id_to_class = {i: class_names[i] for i in range(len(class_names))}

    device = (
        torch.device("cuda:1") if torch.cuda.is_available() else torch.device("cpu")
    )
    args = load_params("MPTR_AutoT_seed0_args.json")
    model_path = f"IKMLab/MPTR_AutoT"
    tokenizer = AutoTokenizer.from_pretrained(model_path)

    if prompt_FT:
        # Prompt tuning
        label_words = get_label_words(list(classes.keys()), args.use_multi_label_words)

        if args.use_multi_label_words:
            label_word_ids = []
            for l in label_words:
                one_label_ids = [tokenizer.convert_tokens_to_ids(word) for word in l]
                label_word_ids.append(one_label_ids)
        else:
            label_word_ids = (
                torch.tensor([tokenizer.convert_tokens_to_ids(l) for l in label_words])
                .long()
                .to(device)
            )
        model = BertForPromptFinetuning.from_pretrained(
            model_path,
            use_multi_label_words=args.use_multi_label_words,
        )
        model.label_word_ids = label_word_ids

    result_path = f"results/predict"

    training_args = TrainingArguments(
        output_dir=result_path,
        learning_rate=args.lr,
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=1,
        num_train_epochs=args.num_epochs,
        weight_decay=0.01,
        warmup_ratio=args.warmup_ratio,
        seed=args.seed,
        evaluation_strategy="steps",
        logging_steps=100,  # same as eval_steps
        save_strategy="steps",
        save_steps=100,
        save_total_limit=1,
        load_best_model_at_end=True,
        metric_for_best_model=f"eval_{args.best_metric}",
    )
    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=None,
        eval_dataset=None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics(
            threshold=args.t,
            classes=classes,
            p_tuning=prompt_FT,
        ),
    )

    testset = InferenceDataset(
        Liver_CT_Report,
        tokenizer,
        args.max_seq_len,
        template=args.template,
        prompt=args.prompt,
    )
    result = trainer.predict(testset)
    predictions = (result.predictions[0] >= args.t) * 1
    positive_idx = np.where(predictions == 1)[0]
    if len(positive_idx) == 0:
        return "No positive findings."

    return [id_to_class[i] for i in positive_idx]


# test = "Two small 0.6-cm and 1.4-cm densely packed lipiodol puddles in S7 without identifiable viable tumor, suggestive of good response to previous TACE without viability."
# result = greet(test)
# Define your custom CSS
custom_css = """
    div.svelte-1viwdyg {
        text-align: left;
    }
"""


iface = gr.Interface(
    fn=greet,
    inputs="text",
    outputs="text",
    title="MPTR+AutoT",
    description="You can try the three examples provided below (same as the ones in our paper) or other liver CT reports.",
    examples=[
        "Comparsion: CT study on Technique: Triphasic CT study of the liver with 5 mm spiral contiguous helical slice was obtained through the abdomen following the uneventful administration of 100 cc IV contrast enhanced CT in arterial phase, portal phase and delay scan shows: > Mild undulated surface of shrunken liver, borderline splenomegaly, suggesting liver cirrhosis and portal hypertension.  > Foci arterial blushes in S7 and S8 near liver dome, favor of AP shunts. No obvious abnormal tumor blushes nor enhancement nodule to be noted in the other part of liver. A tiny hepatic cyst in S7/8> Patency of the SMV, splenic vein, portal and hepatic veins. No obvious biliary tree dilatation. > No obvious ascites nor mesenteric massThe spleen, pancreas, bilateral adrenals and both kidneys are unremarkable. > No definite enlarged retroperitoneal LNs in the abdomen and pelvis. > Some divertiucla in sigmoid colon. Juxtampullar divertiuclum at 2nd portion of duodednum. Patent of the bowel loop without obvious eccentric mass> No obvious pulmonary nodule in bilateral lower lungs. > Degenerative change of the L-spine. Others: Atheroscleroctic change of aorta and its major branches. IMP:No obvious hypervascular nodule nor HCC in both lobes liver. Foci AP shunts in S7 and S8 near liver dome. Mild liver cirrhosis and portal hypertension.",
        "A 74 Y/O male; Clinical Information:Portal Hypertension: nil, TACE on  (AFP; 7.0), TACE on , TACE on  (AFP; 7.6), .CT scan of liver for F/U a patient of HCC post operation and TACE for recurrent HCCs was done by using triphasic study without and with bolus IV non-ionic contrast enhancement showed: 1. S/P heaptectomy of left lobe and partial right segmentectomy of S6 as well as cholecystectomy. 2. No imaging evidence of cirrhosis of liver. 3. Two small  0.6-cm and 1.4-cm densely packed lipiodol puddles in S7 without identifiable viable tumor, suggestive of good response to previous TACE without viability. 4. A 0.5-cm nodular enhancement in S8 is noted on the arterial phase image (se 9, im 5). 5. Multiple hepatic cysts are noted, stationary. 6. The portal and hepatic venous system are patent. No biliary tree dilatation. 7. No remarkable finding of the spleen, pancreas, both kidneys and adrenal glands. 8. No evidence of enlarged lymph node is found in the perigastric area, hepatoduondenal ligament, para-aortic area, pelvis and inguina. 9. Grossly, no abnormality is found in the GI tract. Clear mesentery and omentum. No ascites. 10. Normal contour, capacity and wall thickness of urinary bladder. Normal size and contour of seminal vesicle and prostatic gland. 11. Chronic fibrotic change with reticulonodular infiltrations are noted in bilateral upper lungs, old pulmonary TB should be susptected. Mild degree reticulonodular infiltrations in RML. 12. No enlarged lymph node or tumor mass is found in the mediastinum. 13. Grossly, no destructive bony lesion or abnormal bone density. IMP: Two small HCCs in S7 post successful TACE without viability. Suspicion of a 0.5-cm newly found HCC in S8. F/U dynamic study 3-6 months later is indicated.",
        "A 74 Y/O male; Clinical Information: umbilical painful mass 4*4cm, nature? CT of abdomen without & with contrast enhancement shows: 1. S/P radical prostatectomy; no gross local recurrence but a 3.6-cm mass in the anterior abdominal wall, R/O metastasis; suggest clinical correlation 2. Mild bilateral renal atrophy; right renal cyst, 3 cm; remarkable fatty liver with GB stones; no gross dilatation in the biliary tree 3. No remarkable finding in the pancreas, spleen, and adrenal glands 4. No enlarged lymph nodes at the paraaortic and iliac chain areas. 5. No ascites; clear bilateral basal lungs; no gross bony metastasisIMP: S/P prostatectomy;  R/O metastasis in the anterior abdominal wall",
    ],
    css=custom_css,
    cache_examples=False,
)
iface.launch(share=True)