feat: first commit

app.py

@@ -0,0 +1,152 @@
+from src.modeling_t5 import T5ForSequenceClassification
+import selfies as sf
+import pandas as pd
+from transformers import AutoTokenizer, pipeline
+from chemistry_adapters.amino_acids import AminoAcidAdapter
+from tqdm import tqdm
+import gradio as gr
+
+
+class xBitterT5_predictor:
+    def __init__(
+        self,
+        xBitterT5_640_ckpt="cbbl-skku-org/xBitterT5-640",
+        xBitterT5_720_ckpt="cbbl-skku-org/xBitterT5-720",
+        device="cpu",
+    ):
+        self.xBitterT5_640_ckpt = xBitterT5_640_ckpt
+        self.xBitterT5_720_ckpt = xBitterT5_720_ckpt
+        self.device = device
+
+        self.tokenizer = AutoTokenizer.from_pretrained(xBitterT5_640_ckpt)
+        self.xBitterT5_640 = self.load_model(xBitterT5_640_ckpt)
+        self.xBitterT5_720 = self.load_model(xBitterT5_720_ckpt)
+
+        self.classifier_640 = pipeline(
+            "text-classification",
+            model=self.xBitterT5_640,
+            tokenizer=self.tokenizer,
+            device=self.device,
+        )
+        self.classifier_720 = pipeline(
+            "text-classification",
+            model=self.xBitterT5_720,
+            tokenizer=self.tokenizer,
+            device=self.device,
+        )
+
+    def load_model(self, ckpt):
+        model = T5ForSequenceClassification.from_pretrained(ckpt)
+        model.eval()
+        model.to(self.device)
+        return model
+
+    def convert_sequence_to_smiles(self, sequence):
+        adapter = AminoAcidAdapter()
+        return adapter.convert_amino_acid_sequence_to_smiles(sequence)
+
+    def conver_smiles_to_selfies(self, smiles):
+        return sf.encoder(smiles)
+
+    def predict(
+        self,
+        input_dict,
+        model_type="xBitterT5-720",
+        batch_size=4,
+    ):
+        assert model_type in ["xBitterT5-640", "xBitterT5-720"]
+        df = pd.DataFrame(
+            {"id": list(input_dict.keys()), "sequence": list(input_dict.values())}
+        )
+
+        df["smiles"] = df.apply(
+            lambda row: self.convert_sequence_to_smiles(row["sequence"]),
+            axis=1,
+        )
+        df["selfies"] = df.apply(
+            lambda row: self.conver_smiles_to_selfies(row["smiles"]),
+            axis=1,
+        )
+
+        df["sequence"] = df.apply(
+            lambda row: "<bop>"
+            + "".join("<p>" + aa for aa in row["sequence"])
+            + "<eop>",
+            axis=1,
+        )
+        df["selfies"] = df.apply(lambda row: "<bom>" + row["selfies"] + "<eom>", axis=1)
+        df["text"] = df["sequence"] + df["selfies"]
+
+        text_inputs = df["text"].tolist()
+
+        if model_type == "xBitterT5-640":
+            classifier = self.classifier_640
+        else:
+            classifier = self.classifier_720
+
+        result = []
+        for i in tqdm(range(0, len(text_inputs), batch_size)):
+            batch = text_inputs[i : i + batch_size]
+            result.extend(classifier(batch))
+
+        y_pred, y_prob = [], []
+        for pred in result:
+            if pred["label"] == "bitter":
+                y_prob.append(pred["score"])
+                y_pred.append(1)
+            else:
+                y_prob.append(1 - pred["score"])
+                y_pred.append(0)
+
+        return {i: [y_prob[j], y_pred[j]] for j, i in enumerate(df["id"].tolist())}
+
+
+predictor = xBitterT5_predictor()
+
+
+def process_fasta(fasta_text):
+    """
+    Processes the input FASTA format text into a dictionary {id: sequence}.
+    """
+    fasta_dict = {}
+    current_id = None
+    current_sequence = []
+
+    for line in fasta_text.strip().split("\n"):
+        line = line.strip()
+        if line.startswith(">"):  # Header line
+            if current_id:
+                fasta_dict[current_id] = "".join(current_sequence)
+            current_id = line[1:]  # Remove '>'
+            current_sequence = []
+        else:
+            current_sequence.append(line)
+
+    # Add the last sequence
+    if current_id:
+        fasta_dict[current_id] = "".join(current_sequence)
+
+    return fasta_dict
+
+
+# Create a Gradio interface
+def gradio_process_fasta(fasta_text):
+    """
+    Wrapper for Gradio to process the FASTA text.
+    """
+    fasta_dict = process_fasta(fasta_text)
+    result = predictor.predict(fasta_dict)
+    return result
+
+
+interface = gr.Interface(
+    fn=gradio_process_fasta,
+    inputs=gr.Textbox(
+        label="Enter FASTA format text", lines=10, placeholder=">id1\nATGC\n>id2\nCGTA"
+    ),
+    outputs=gr.JSON(label="Processed FASTA Dictionary with Probabilities and Classes"),
+    title="FASTA to Dictionary with Probabilities and Classes",
+    description=("Enter a FASTA-formatted text"),
+)
+# Launch the Gradio app
+interface.launch()

requirements.txt

@@ -0,0 +1,97 @@
+aiohappyeyeballs==2.4.4
+aiohttp==3.11.11
+aiosignal==1.3.2
+async-timeout==5.0.1
+attrs==24.3.0
+backcall==0.2.0
+captum==0.7.0
+certifi==2024.12.14
+charset-normalizer==3.4.1
+chemistry_adapters==0.0.2
+contourpy==1.3.0
+cycler==0.12.1
+datasets==3.2.0
+decorator==5.1.1
+dill==0.3.8
+evaluate==0.4.3
+filelock==3.16.1
+fonttools==4.55.3
+frozenlist==1.5.0
+fsspec==2024.9.0
+huggingface-hub==0.27.1
+idna==3.10
+importlib_resources==6.5.2
+ipython==7.34.0
+jedi==0.19.2
+Jinja2==3.1.5
+joblib==1.4.2
+kiwisolver==1.4.7
+MarkupSafe==3.0.2
+matplotlib==3.9.4
+matplotlib-inline==0.1.7
+mpmath==1.3.0
+multidict==6.1.0
+multiprocess==0.70.16
+networkx==3.2.1
+numpy==2.0.2
+nvidia-cublas-cu11==11.11.3.6
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu11==11.8.87
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu11==11.8.89
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu11==11.8.89
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu11==9.1.0.70
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu11==10.9.0.58
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu11==10.3.0.86
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu11==11.4.1.48
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu11==11.7.5.86
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu11==2.21.5
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu11==11.8.86
+nvidia-nvtx-cu12==12.1.105
+packaging==24.2
+pandas==2.2.3
+parso==0.8.4
+pexpect==4.9.0
+pickleshare==0.7.5
+pillow==11.1.0
+prompt_toolkit==3.0.50
+propcache==0.2.1
+ptyprocess==0.7.0
+pyarrow==19.0.0
+Pygments==2.19.1
+pyparsing==3.2.1
+python-dateutil==2.9.0.post0
+pytz==2024.2
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.3
+safetensors==0.5.2
+scikit-learn==1.6.1
+scipy==1.13.1
+selfies==2.2.0
+six==1.17.0
+sympy==1.13.1
+threadpoolctl==3.5.0
+tokenizers==0.21.0
+torch==2.5.1+cu121
+tqdm==4.67.1
+traitlets==5.14.3
+transformers==4.48.1
+transformers-interpret==0.10.0
+triton==3.1.0
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.3.0
+wcwidth==0.2.13
+xxhash==3.5.0
+yarl==1.18.3
+zipp==3.21.0

src/data.py

@@ -0,0 +1,101 @@
+from datasets import Dataset, DatasetDict
+import pandas as pd
+import glob
+import os
+import numpy as np
+
+
+def create_dataset_from_dataframe(
+    dataframe_path, pretrained_name, chosen_features=None
+):
+    dataframe = pd.read_csv(dataframe_path, usecols=["label"] + chosen_features)
+    rows_with_nan = dataframe[chosen_features].isna().any(axis=1)
+    dataframe = dataframe[np.logical_not(rows_with_nan)]
+    if len(chosen_features) > 1:
+        for feature in chosen_features:
+            if feature == "selfies":
+                dataframe[feature] = dataframe.apply(
+                    lambda row: "<bom>" + row[feature] + "<eom>", axis=1
+                )
+            elif feature == "sequence":
+                dataframe[feature] = dataframe.apply(
+                    lambda row: "<bop>"
+                    + "".join("<p>" + aa for aa in row[feature])
+                    + "<eop>",
+                    axis=1,
+                )
+
+            dataframe["text"] = dataframe.apply(
+                lambda row: "".join([f"{row[feature]}" for feature in chosen_features]),
+                axis=1,
+            )
+
+    elif len(chosen_features) == 1:
+        chosen_feature = chosen_features[0]
+        if chosen_feature == "selfies":
+            dataframe["text"] = dataframe.apply(
+                lambda row: "<bom>" + row[chosen_feature] + "<eom>", axis=1
+            )
+        elif chosen_feature == "smiles":
+            dataframe["text"] = dataframe[chosen_feature]
+        elif chosen_feature == "sequence":
+            if "biot5" in pretrained_name:
+                dataframe["text"] = dataframe.apply(
+                    lambda row: "<bop>"
+                    + "".join("<p>" + aa for aa in row[chosen_feature])
+                    + "<eop>",
+                    axis=1,
+                )
+            else:
+                dataframe["text"] = dataframe.apply(
+                    lambda row: " ".join(row[chosen_feature]),
+                    axis=1,
+                )
+    dataframe.drop(columns=chosen_features, inplace=True)
+    dataset = Dataset.from_pandas(dataframe)
+    return dataset
+
+
+def create_and_save_datadict(train, val, test, save_path):
+    if val is None:
+        dataset_dict = DatasetDict({"train": train, "test": test})
+        dataset_dict.save_to_disk(save_path)
+        return dataset_dict
+    dataset_dict = DatasetDict({"train": train, "val": val, "test": test})
+    dataset_dict.save_to_disk(save_path)
+    return dataset_dict
+
+
+def prepare_dataset(args):
+    fold_folders = glob.glob(args.data_folder + "/fold_*/")
+    for fold_folder in fold_folders:
+        train_path = os.path.join(fold_folder, "train.csv")
+        val_path = os.path.join(fold_folder, "val.csv")
+        test_path = os.path.join(fold_folder, "test.csv")
+
+        train = create_dataset_from_dataframe(
+            train_path, args.pretrained_name, args.chosen_features
+        )
+        val = create_dataset_from_dataframe(
+            val_path, args.pretrained_name, args.chosen_features
+        )
+        test = create_dataset_from_dataframe(
+            test_path, args.pretrained_name, args.chosen_features
+        )
+        folder_name = f"dataset_{'_'.join(args.chosen_features)}_{args.pretrained_name.split('/')[-1].replace('-', '_')}"
+        save_path = os.path.join(fold_folder, folder_name)
+        create_and_save_datadict(train, val, test, save_path)
+
+    train_path = os.path.join(args.data_folder, "train.csv")
+    test_path = os.path.join(args.data_folder, "test.csv")
+    train = create_dataset_from_dataframe(
+        train_path, args.pretrained_name, args.chosen_features
+    )
+    test = create_dataset_from_dataframe(
+        test_path, args.pretrained_name, args.chosen_features
+    )
+    save_path = os.path.join(
+        args.data_folder,
+        f"dataset_{'_'.join(args.chosen_features)}_{args.pretrained_name.split('/')[-1].replace('-', '_')}",
+    )
+    create_and_save_datadict(train, None, test, save_path)

src/explainer.py

@@ -0,0 +1,32 @@
+from transformers_interpret import SequenceClassificationExplainer
+from typing import List, Tuple, Union
+import torch
+
+
+class xBitterT5_explainer(SequenceClassificationExplainer):
+    def _make_input_reference_pair(
+        self, text: Union[List, str]
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        if isinstance(text, list):
+            raise NotImplementedError("Lists of text are not currently supported.")
+
+        text_ids = self.encode(text)
+        input_ids = self.tokenizer.encode(text, add_special_tokens=True)
+
+        # if no special tokens were added
+        if len(text_ids) == len(input_ids):
+            ref_input_ids = [self.ref_token_id] * len(text_ids)
+        else:
+            ref_input_ids = (
+                [self.cls_token_id]
+                + [self.ref_token_id] * len(text_ids)
+                + [self.sep_token_id]
+            )
+
+        # Use this because pretrained BioT5 plus does not have cls_token_id
+        ref_input_ids = [self.ref_token_id] * len(text_ids) + [self.sep_token_id]
+        return (
+            torch.tensor([input_ids], device=self.device),
+            torch.tensor([ref_input_ids], device=self.device),
+            len(text_ids),
+        )

src/model.py

@@ -0,0 +1,58 @@
+# from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
+from transformers import AutoTokenizer, T5Tokenizer, AutoConfig
+from src.modeling_t5 import T5ForSequenceClassification
+
+
+def prepare_tokenizer(args):
+    try:
+        try:
+            return AutoTokenizer.from_pretrained(args.pretrained_name)
+        except Exception as e:
+            print(f"Error: {e}")
+            return T5Tokenizer.from_pretrained(
+                args.pretrained_name,
+                do_lower_case=False,
+            )
+    except Exception as e:
+        print(f"Error: {e}")
+        return T5Tokenizer.from_pretrained(args.pretrained_name)
+
+
+def check_unfreeze_layer(name, trainable_layers):
+    flag = False
+    for layer in trainable_layers:
+        if name.startswith(f"transformer.decoder.block.{layer}"):
+            flag = True
+            break
+    return flag
+
+
+def prepare_model(args):
+    id2lable = {0: "non-bitter", 1: "bitter"}
+    label2id = {"non-bitter": 0, "bitter": 1}
+    config = AutoConfig.from_pretrained(
+        args.pretrained_name,
+        cache_dir=args.cache_dir,
+        num_labels=2,
+        id2label=id2lable,
+        label2id=label2id,
+    )
+    config.dropout_rate = args.dropout
+    config.classifier_dropout = args.dropout
+    config.problem_type = "single_label_classification"
+
+    model = T5ForSequenceClassification.from_pretrained(
+        args.pretrained_name,
+        cache_dir=args.cache_dir,
+        config=config,
+    )
+    model.to(args.accelerator)
+    for name, param in model.named_parameters():
+        if name.startswith("classification_head") or check_unfreeze_layer(
+            name, args.trainable_layers
+        ):
+            param.requires_grad = True
+        else:
+            param.requires_grad = False
+
+    return model

src/utils.py

@@ -0,0 +1,43 @@
+import evaluate
+import numpy as np
+from datetime import datetime
+from zoneinfo import ZoneInfo
+from torch.nn.functional import softmax
+from torch import tensor
+from sklearn.metrics import confusion_matrix, roc_curve, auc
+
+
+bitter_metrics = evaluate.combine(
+    ["accuracy", "f1", "precision", "recall", "matthews_correlation"]
+)
+
+
+def compute_metrics(eval_pred):
+    predictions, labels = eval_pred
+    preds = np.argmax(predictions[0], axis=1)
+    prediction_scores = softmax(tensor(predictions[0]), dim=-1)
+    prediction_scores = prediction_scores[:, 1].cpu().numpy()
+
+    metrics = bitter_metrics.compute(predictions=preds, references=labels)
+    tn, fp, fn, tp = confusion_matrix(labels, preds).ravel()
+    specificity = tn / (tn + fp)
+    metrics.update(
+        {
+            "eval_specificity": specificity,
+            "eval_tn": tn,
+            "eval_fp": fp,
+            "eval_fn": fn,
+            "eval_tp": tp,
+        }
+    )
+
+    fpr2, tpr2, _ = roc_curve(labels, prediction_scores, pos_label=1)
+    auc2 = auc(fpr2, tpr2)
+    metrics.update({"eval_auc": auc2})
+
+    metrics = dict(sorted(metrics.items()))
+    return metrics
+
+
+def get_time_string():
+    return datetime.now(tz=ZoneInfo("Asia/Seoul")).strftime("%Y_%m_%d__%H_%M_%S")