Add application files

app.py

@@ -0,0 +1,127 @@
+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(input_text):
+    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("args.json")
+    model_path = f"IKMLab/MPTR"
+    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(
+        input_text,
+        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)
+
+
+iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+iface.launch()

args.json

@@ -0,0 +1,33 @@
+{
+    "template": "*cls**sent_0*_[PROMPT]*mask*.*sep+*",
+    "prompt": "The_report_is_related_to",
+    "num_labels": 7,
+    "report_filter": "full",
+    "max_seq_len": 512,
+    "batch_size": 2,
+    "model_name": "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract",
+    "gpu_id": "1",
+    "cls_mode": "multi_label",
+    "k": null,
+    "t": 0.2,
+    "db_date": "20230606_new",
+    "best_metric": "loss",
+    "exp_tag": "",
+    "num_exps": 5,
+    "do_train": true,
+    "num_epochs": 30,
+    "do_predict": true,
+    "seed": 42,
+    "lr": 3e-05,
+    "warmup_ratio": 0.0,
+    "data_type": "train_32",
+    "save_conf_matrix": false,
+    "use_multi_label_words": true,
+    "allow_multi_label_tokens": false,
+    "verbalizer_name": "",
+    "enable_emboliz": false,
+    "enable_rfa": false,
+    "enable_tace": false,
+    "enable_lobectomy": false,
+    "save_checkpoints": true
+}

class_names.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d8088c6c9d790808303e0a3e9b122c9d9103a4e2c30694f4d5e3351d2c25872
+size 110

prompt_dataset.py

@@ -0,0 +1,169 @@
+import torch
+import pandas as pd
+
+
+def get_prompt_length(tokenizer, prompt):
+    return len(tokenizer.encode(prompt))
+
+
+def tokenize_multipart_input(
+    tokenizer,
+    input_text_list: list,
+    max_seq_len: int,
+    template=None,
+    prompt=None,
+):
+    """This function is an adaptation of the `tokenize_multipart_input` found in princeton-nlp's repository
+    at https://github.com/princeton-nlp/LM-BFF/blob/main/src/dataset.py.
+
+    Modifications include:
+    - Extension of automatic prompt generation for multi-label classification.
+    - Removal of parameters like `first_sent_limit`, `other_sent_limit`, `gpt3`, `truncate_head`, and `support_labels`.
+    - Optimization of the code flow.
+
+    Args:
+        tokenizer: a pre-trained tokenizer from Hugging Face Transformers
+        input_text_list (list): documents ready for tokenization.
+        max_seq_len (int): max sequence length after adding the prompt along with special tokens from BERT.
+        template (str, optional): placeholder for the prompt.
+        prompt (str, optional): the prompt we use for input text.
+    """
+
+    def enc(text):
+        return tokenizer.encode(text, add_special_tokens=False)
+
+    input_ids = []
+    attention_mask = []
+    token_type_ids = []  # Only for BERT
+    mask_pos = None  # Position of the mask token
+
+    if prompt:
+        special_token_mapping = {
+            "cls": tokenizer.cls_token_id,
+            "mask": tokenizer.mask_token_id,
+            "sep": tokenizer.sep_token_id,
+            "sep+": tokenizer.sep_token_id,
+        }
+        # Get variable list in the template
+        if prompt != "auto":
+            template = template.replace("[PROMPT]", prompt)
+        template_list = template.split("*")
+        if prompt == "auto":
+            # find cls place
+            cls_pos = template_list.index("cls")
+            if template_list[cls_pos + 1] == "":
+                # For these kinds of cases: *cls**sent_0*_Liver*mask*.*sep+*
+                # Prompt is next to sent_0.
+                prompt = template_list[cls_pos + 3]
+            elif template_list[cls_pos + 1] != "" and (
+                template_list[cls_pos + 1].startswith("_")
+            ):
+                # For these kinds of cases: *cls*_Liver*mask*.*+sent_0**sep+*
+                # Prompt is next to cls.
+                prompt = template_list[cls_pos + 1]
+            if prompt.startswith("_"):
+                prompt = prompt[1:]
+        segment_id = 0
+
+        for part in template_list:
+            new_tokens = []
+            segment_plus_1_flag = False
+            if part in special_token_mapping:
+                new_tokens.append(special_token_mapping[part])
+                if part == "sep+":
+                    segment_plus_1_flag = True
+            elif part[:5] == "sent_" or part[:6] == "+sent_":
+                sent_id = int(part.split("_")[1])
+                max_len = max_seq_len - 3 - get_prompt_length(tokenizer, prompt)
+                # Tokenize and truncate to max_seq_len
+                tokens = enc(input_text_list[sent_id])[-max_len:]
+                new_tokens += tokens
+            else:
+                # Just natural language prompt
+                part = part.replace("_", " ")
+                # handle special case when T5 tokenizer might add an extra space
+                if len(part) == 1:
+                    new_tokens.append(tokenizer.convert_tokens_to_ids(part))
+                else:
+                    new_tokens += enc(part)
+
+            input_ids += new_tokens
+            attention_mask += [1 for i in range(len(new_tokens))]
+            token_type_ids += [segment_id for i in range(len(new_tokens))]
+
+            if segment_plus_1_flag:
+                segment_id += 1
+
+        mask_pos = [input_ids.index(tokenizer.mask_token_id)]
+        # Make sure that the masked position is inside the max_length
+        assert mask_pos[0] < max_seq_len
+
+    else:
+        input_ids = [tokenizer.cls_token_id]
+        attention_mask = [1]
+        token_type_ids = [0]
+        max_len = max_seq_len - 2
+
+        for sent_id, input_text in enumerate(input_text_list):
+            if input_text is None:
+                # Do not have text_b
+                continue
+            if pd.isna(input_text) or input_text is None:
+                # Empty input
+                input_text = ""
+            input_tokens = enc(input_text)[:max_len] + [tokenizer.sep_token_id]
+            input_ids += input_tokens
+            attention_mask += [1 for i in range(len(input_tokens))]
+            token_type_ids += [sent_id for i in range(len(input_tokens))]
+
+    return input_ids, attention_mask, token_type_ids, mask_pos
+
+
+class InferenceDataset(torch.utils.data.Dataset):
+    """
+    A class for creating the CGMH dataset in PyTorch.
+    Currently, this class supports:
+    (1) Few-shot data (e.g., train_size=16)
+    (2) Small-size data (e.g., train_size>100)
+    ---
+    Attributes
+        data (pd.DataFrame): the CGMH dataset
+        tokenizer: a pre-trained HuggingFace tokenizer
+        max_seq_len (int): maximum length for a sequence
+        template (_type_, optional): template for the model. Defaults to None.
+        prompt (_type_, optional): prompt for the model. Defaults to None.
+    """
+
+    def __init__(
+        self,
+        input_text: str,
+        tokenizer,
+        max_seq_len: int,
+        template=None,
+        prompt=None,
+    ):
+        self.doc = input_text
+        self.template = template
+        self.prompt = prompt
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+
+    def __getitem__(self, idx):
+        input_ids, attn_mask, segs, mask_pos = tokenize_multipart_input(
+            tokenizer=self.tokenizer,
+            input_text_list=[self.doc],
+            template=self.template,
+            prompt=self.prompt,
+            max_seq_len=self.max_seq_len,
+        )
+        item = {
+            "input_ids": input_ids,
+            "token_type_ids": segs,
+            "attention_mask": attn_mask,
+        }
+        if self.prompt:
+            item["mask_pos"] = mask_pos
+        return item
+
+    def __len__(self):
+        return 1

prompt_model_factory.py

@@ -0,0 +1,88 @@
+from typing import Optional
+from transformers import BertModel
+from transformers.models.bert.modeling_bert import (
+    BertPreTrainedModel,
+    BertOnlyMLMHead,
+)
+import torch
+
+
+class BertForPromptFinetuning(BertPreTrainedModel):
+    def __init__(self, config, use_multi_label_words: bool = False):
+        super().__init__(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+        # Initialize weights and apply final processing
+        self.init_weights()
+
+        self.label_word_ids = None
+        self.use_multi_label_words = use_multi_label_words
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        mask_pos: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = False,
+        output_attentions: Optional[bool] = False,
+    ):
+        if mask_pos is not None:
+            mask_pos = mask_pos.squeeze()
+        elif mask_pos is None:
+            raise ValueError("`mask_pos` should be assigned!")
+
+        # Encode everything
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+
+        # Get <mask> token representation
+        sequence_output = outputs[0]
+        sequence_mask_output = sequence_output[
+            torch.arange(sequence_output.size(0)), mask_pos
+        ]
+
+        # Logits over vocabulary tokens
+        # prediction_mask_scores.shape: [batch_size, vocab_size]
+        prediction_mask_scores = self.cls(sequence_mask_output)
+
+        # Return logits for each label
+        logits = []
+        if self.use_multi_label_words:
+            for label_id in self.label_word_ids:
+                one_label_logits = []
+                # multiple ids in one label_id
+                for id in label_id:
+                    one_label_word_logits = prediction_mask_scores[:, id]
+                    one_label_logits.append(one_label_word_logits.unsqueeze(-1))
+                # one_label_logits: (bs, num_label_words)
+                one_label_logits = torch.cat(one_label_logits, -1)
+                # Get the max logits to choose the label word
+                logits.append(torch.max(one_label_logits, dim=1, keepdim=True)[0])
+
+        else:
+            for label_id in range(len(self.label_word_ids)):
+                logits.append(
+                    prediction_mask_scores[:, self.label_word_ids[label_id]].unsqueeze(
+                        -1
+                    )
+                )
+
+        # logits.shape: [batch_size, num_classes]
+        logits = torch.sigmoid(torch.cat(logits, -1))
+
+        loss = None
+        if labels is not None:
+            loss_fct = torch.nn.BCELoss()
+            loss = loss_fct(logits, labels.float())
+
+        output = (logits, outputs.hidden_states) if output_hidden_states else (logits,)
+        output = (output + (outputs.attentions)) if output_attentions else output
+
+        return ((loss,) + output) if loss is not None else output

utils.py

@@ -0,0 +1,138 @@
+from pathlib import Path
+import torch
+import os
+import random
+import argparse
+import json
+import pandas as pd
+import numpy as np
+from sklearn.metrics import precision_recall_fscore_support
+from ast import literal_eval
+
+
+def pred_by_threshold(
+    t: float,
+    y_true: np.array,
+    similarities: np.array,
+    classes: dict,
+):
+    preds = (similarities >= t) * 1
+    sk_results = precision_recall_fscore_support(
+        y_true,
+        preds,
+        # average="samples",  # For calculating sample-wise P and R scores.
+    )
+    outputs = {
+        "f1": np.average(sk_results[2]),
+        "P": np.average(sk_results[0]),
+        "R": np.average(sk_results[1]),
+    }
+    for label_name, idx in classes.items():
+        outputs[f"{label_name}_f1"] = sk_results[2][idx]
+    return outputs
+
+
+def get_avg_length(dataset: torch.utils.data.Dataset):
+    all_lengths = 0
+    data_size = len(dataset)
+    for i in range(data_size):
+        all_lengths += len(dataset[i]["input_ids"])
+    return all_lengths / data_size
+
+
+def load_csv_multi_label(filename: str, col_name: str = "labels") -> pd.DataFrame:
+    """Prevent Pandas from converting lists of int into lists of strings.
+
+    Args:
+        filename (str): path of a csv file
+        col_name (str, optional): column name of lists of int. Defaults to 'labels'.
+
+    Returns:
+        pd.DataFrame: a Pandas dataframe
+    """
+    return pd.read_csv(filename, converters={col_name: literal_eval})
+
+
+def save_logged_results(filename: str, results: dict):
+    try:
+        old_df = pd.read_csv(filename)
+        df = pd.concat([old_df, pd.DataFrame(results)], ignore_index=True)
+    except FileNotFoundError:
+        df = pd.DataFrame(results)
+
+    df.to_csv(filename, index=None)
+
+
+def set_seed(seed):
+    """
+    Args:
+        seed: an integer number to initialize a pseudorandom number generator
+    """
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        # torch.cuda.manual_seed_all(seed)  # if using more than one GPUs
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+def save_baseline_table(
+    y_preds: list,
+    baseline_name: str,
+    baseline_result_file: str = "results/baselines.pkl",
+    all_doc_idx: list = None,
+) -> None:
+    if Path(baseline_result_file).exists():
+        df = pd.read_pickle(baseline_result_file)
+    else:
+        assert all_doc_idx is not None
+        df = pd.DataFrame({"doc_idx": all_doc_idx})
+
+    df[baseline_name] = y_preds
+    df.to_pickle(baseline_result_file)
+
+
+def load_params(path_of_params):
+    with open(path_of_params, "r") as f:
+        params = json.load(f)
+    return argparse.Namespace(**params)
+
+
+def get_label_words(classes: list, use_multi_label_words=False) -> list:
+    mapping = {
+        "cyst": "cyst",
+        "HCC": "hcc",  # hepatoma
+        "cirrhosis": "cirrhosis",
+        "post-treatment": "posttreatment",
+        "steatosis": "steatosis",
+        "metastasis": "metastasis",
+        "hemangioma": "hemangioma",
+    }
+    if use_multi_label_words:
+        mapping = {
+            "cyst": ["cyst"],
+            "HCC": ["hcc", "hepatoma"],  # hepatoma
+            "cirrhosis": ["cirrhosis"],
+            "post-treatment": ["posttreatment"],
+            "steatosis": ["steatosis", "steatohepatitis"],
+            "metastasis": ["metastasis"],
+            "hemangioma": ["hemangioma"],
+        }
+
+    label_words = [mapping[c] for c in classes]
+    return label_words
+
+
+def seed_mapper(data_type: str) -> list:
+    mapping = {
+        "train_8": [2, 4, 7, 11, 21, 23, 24, 36, 44, 128],
+        "train_32": [0, 1, 3, 7, 10],
+    }
+    if data_type in mapping:
+        return mapping[data_type]
+    else:
+        raise NotImplementedError