app

.ipynb_checkpoints/app-checkpoint.py

@@ -33,6 +33,9 @@ sys.path.append( '../')
 import tokenizer_sl, datamodule_finetune_sl, model_finetune_sl, chemllama_mtr, utils_sl
 import auto_evaluator_sl
 
+from torch.utils.data import Dataset, DataLoader
+from transformers import DataCollatorWithPadding
+
 torch.manual_seed(1004)
 np.random.seed(1004)
 
@@ -48,26 +51,143 @@ with open(file_path, 'w') as file:
 ###
 # solute_or_solvent = 'solute'
 solute_or_solvent = st.selectbox('Solute or Solvent', ['Solute', 'Solvent'])
-ver_ft = 0 # version control for FT model & evaluation data # Or it will overwrite the models and results
-batch_size_pair = [64, 64] if solute_or_solvent == 'Solute' else [10, 10] # [train, valid(test)] 
-# since 'solute' has very small dataset. So I thinl 10 for train and 10 for valid(test) should be the maximum values.
-lr = 0.0001 
-epochs = 7
-use_freeze = False  # Freeze the model or not # False measn not freezing
-overwrite_level_2 = True
-###
-max_seq_length = 512
-tokenizer = tokenizer_sl.fn_load_tokenizer_llama(
-    max_seq_length=max_seq_length,
-)
-max_length = max_seq_length
-num_workers = 2
+
+
+class ChemLlama(nn.Module):
+    def __init__(
+        self,
+        max_position_embeddings=512,
+        vocab_size=591,
+        pad_token_id=0,
+        bos_token_id=12,
+        eos_token_id=13,
+        hidden_size=768,
+        intermediate_size=768,
+        num_labels=105,
+        attention_dropout=0.144,
+        num_hidden_layers=7,
+        num_attention_heads=8,
+        learning_rate=0.0001,
+    ):
+        super(ChemLlama, self).__init__()
+        
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_labels = num_labels
+        self.vocab_size = vocab_size
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.attention_dropout = attention_dropout
+        self.max_position_embeddings = max_position_embeddings
+
+        self.mae = torchmetrics.MeanAbsoluteError()
+        self.mse = torchmetrics.MeanSquaredError()
+
+        self.config_llama = LlamaConfig(
+            max_position_embeddings=self.max_position_embeddings,
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            intermediate_size=self.intermediate_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            attention_dropout=self.attention_dropout,
+            pad_token_id=self.pad_token_id,
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+        )
+
+        self.loss_fn = nn.L1Loss()
+
+        self.llama = LlamaModel(self.config_llama)
+        self.gelu = nn.GELU()
+        self.score = nn.Linear(self.hidden_size, self.num_labels)
+
+    def forward(self, input_ids, attention_mask, labels=None):
+
+        transformer_outputs = self.llama(
+            input_ids=input_ids, attention_mask=attention_mask
+        )
+
+        hidden_states = transformer_outputs[0]
+        hidden_states = self.gelu(hidden_states)
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config_llama.pad_token_id is None and batch_size != 1:
+            raise ValueError(
+                "Cannot handle batch sizes > 1 if no padding token is defined."
+            )
+        if self.config_llama.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+                sequence_lengths = (
+                    torch.eq(input_ids, self.config_llama.pad_token_id).int().argmax(-1)
+                    - 1
+                )
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+        # raise ValueError(len(sequence_lengths), sequence_lengths)
+
+        pooled_logits = logits[
+            torch.arange(batch_size, device=logits.device), sequence_lengths
+        ]
+        return pooled_logits
+
+
+chemllama_mtr = ChemLlama()
+
+class ChemLlama_FT(nn.Module):
+    def __init__(
+        self,
+        model_mtr,
+        linear_param:int=64,
+        use_freeze:bool=True,
+        *args, **kwargs
+    ):
+        super(CustomFinetuneModel, self).__init__()
+        # self.save_hyperparameters()
+
+        self.model_mtr = model_mtr
+        if use_freeze:
+            self.model_mtr.freeze()
+            # for name, param in model_mtr.named_parameters():
+            #     param.requires_grad = False
+            #     print(name, param.requires_grad)
+        
+        self.gelu = nn.GELU()
+        self.linear1 = nn.Linear(self.model_mtr.num_labels, linear_param)
+        self.linear2 = nn.Linear(linear_param, linear_param)
+        self.regression = nn.Linear(linear_param, 5)
+
+        self.loss_fn = nn.L1Loss()
+
+    def forward(self, input_ids, attention_mask, labels=None):
+        x = self.model_mtr(input_ids=input_ids, attention_mask=attention_mask)
+        x = self.gelu(x)
+        x = self.linear1(x)
+        x = self.gelu(x)
+        x = self.linear2(x)
+        x = self.gelu(x)
+        x = self.regression(x)
+        
+        return x
+        
+chemllama_ft = ChemLlama_FT(model_mtr=chemllama_mtr)
+
 
 # I just reused our previous research code with some modifications.
 dir_main = "."
-name_model_mtr = "ChemLlama_Medium_30m_vloss_val_loss=0.029_ep_epoch=04.ckpt" 
-
-dir_model_mtr = f"{dir_main}/SolLlama-mtr/{name_model_mtr}"
 
 max_seq_length = 512
 
@@ -79,75 +199,95 @@ num_workers = 2
 
 ## FT
 
-ver_ft = 0
 dir_model_ft_to_save = f"{dir_main}/SolLlama-mtr"
 # name_model_ft = 'Solvent.pt'
 name_model_ft = f"{solute_or_solvent}.pt"
 
-# Load dataset for finetune
-batch_size_for_train = batch_size_pair[0]
-batch_size_for_valid = batch_size_pair[1]
-
-data_module = datamodule_finetune_sl.CustomFinetuneDataModule(
-    solute_or_solvent=solute_or_solvent,
-    tokenizer=tokenizer,
-    max_seq_length=max_length,
-    batch_size_train=batch_size_for_train,
-    batch_size_valid=batch_size_for_valid,
-    # num_device=int(config.NUM_DEVICE) * config.NUM_WORKERS_MULTIPLIER,
-    num_device=num_workers,
-)
-
-data_module.prepare_data()
-data_module.setup()
-steps_per_epoch = len(data_module.test_dataloader())
-
-# Load model and optimizer for finetune
-learning_rate = lr
-
-
-model_mtr = chemllama_mtr.ChemLlama.load_from_checkpoint(dir_model_mtr)
-
-
-model_ft = model_finetune_sl.CustomFinetuneModel(
-    model_mtr=model_mtr,
-    steps_per_epoch=steps_per_epoch,
-    warmup_epochs=1,
-    max_epochs=epochs,
-    learning_rate=learning_rate,
-    # dataset_dict=dataset_dict,
-    use_freeze=use_freeze,
-)
-
-# 'SolLlama_solute_vloss_val_loss=0.082_ep_epoch=06.pt'
-
-trainer = L.Trainer(
-    default_root_dir=dir_model_ft_to_save,
-    # profiler=profiler,
-    # logger=csv_logger,
-    accelerator='auto',
-    devices='auto',
-    # accelerator='gpu',
-    # devices=[0],
-    min_epochs=1,
-    max_epochs=epochs,
-    precision=32,
-    # callbacks=[checkpoint_callback]
-)
-
-
+# # Load dataset for finetune
+# batch_size_for_train = batch_size_pair[0]
+# batch_size_for_valid = batch_size_pair[1]
+
+# data_module = datamodule_finetune_sl.CustomFinetuneDataModule(
+#     solute_or_solvent=solute_or_solvent,
+#     tokenizer=tokenizer,
+#     max_seq_length=max_length,
+#     batch_size_train=batch_size_for_train,
+#     batch_size_valid=batch_size_for_valid,
+#     # num_device=int(config.NUM_DEVICE) * config.NUM_WORKERS_MULTIPLIER,
+#     num_device=num_workers,
+# )
+
+# data_module.prepare_data()
+# data_module.setup()
+# steps_per_epoch = len(data_module.test_dataloader())
+
+# # Load model and optimizer for finetune
+# learning_rate = lr
+
+
+# model_mtr = chemllama_mtr.ChemLlama.load_from_checkpoint(dir_model_mtr)
+
+
+# model_ft = model_finetune_sl.CustomFinetuneModel(
+#     model_mtr=model_mtr,
+#     steps_per_epoch=steps_per_epoch,
+#     warmup_epochs=1,
+#     max_epochs=epochs,
+#     learning_rate=learning_rate,
+#     # dataset_dict=dataset_dict,
+#     use_freeze=use_freeze,
+# )
+
+# # 'SolLlama_solute_vloss_val_loss=0.082_ep_epoch=06.pt'
+
+# trainer = L.Trainer(
+#     default_root_dir=dir_model_ft_to_save,
+#     # profiler=profiler,
+#     # logger=csv_logger,
+#     accelerator='auto',
+#     devices='auto',
+#     # accelerator='gpu',
+#     # devices=[0],
+#     min_epochs=1,
+#     max_epochs=epochs,
+#     precision=32,
+#     # callbacks=[checkpoint_callback]
+# )
+
+device = 'cpu'
 # Predict
 local_model_ft = utils_sl.load_model_ft_with(
-    class_model_ft=model_ft, 
+    class_model_ft=chemllama_ft, 
     dir_model_ft=dir_model_ft_to_save,
     name_model_ft=name_model_ft
+).to(device)
+
+# result = trainer.predict(local_model_ft, data_module)
+# result_pred = list()
+# result_label = list()
+# for bat in range(len(result)):
+#     result_pred.append(result[bat][0].squeeze())
+#     result_label.append(result[bat][1])
+
+with open('./smiles_str.txt', 'r') as file:
+    smiles_str = file.readline()
+    
+dataset_test = datamodule_finetune_sl.CustomLlamaDatasetAbraham(
+    df=pd.DataFrame([smiles_str]),
+    tokenizer=tokenizer,
+    max_seq_length=max_length
 )
 
-result = trainer.predict(local_model_ft, data_module)
-result_pred = list()
-result_label = list()
-for bat in range(len(result)):
-    result_pred.append(result[bat][0].squeeze())
-    result_label.append(result[bat][1])
+dataloader_test = DataLoader(dataset_test, shuffle=False, collate_fn=DataCollatorWithPadding(tokenizer))
+
+list_predictions = []
+local_model_ft.eval()
+with torch.inference_mode():
+    for i, v_batch in enumerate(dataloader_test):
+        v_input_ids = v_batch['input_ids'].to(device)
+        v_attention_mask = v_batch['attention_mask'].to(device)
+        # v_y_labels = v_batch['labels'].to(device)
+        v_y_logits = local_model_ft(input_ids=v_input_ids, attention_mask=v_attention_mask)
+        list_predictions.append(v_y_logits[0][0].tolist())
 
-st.write(result_pred)
+st.write(list_predictions)

.ipynb_checkpoints/app_old-checkpoint.py

@@ -0,0 +1,153 @@
+
+import streamlit as st
+# from git import Repo
+
+# Repo.clone_from('https://huggingface.co/ttmn/SolLlama-mtr', './SolLlama-mtr')
+import subprocess
+
+def git_clone(repo_url, destination_dir):
+    try:
+        subprocess.run(['git', 'clone', '-v', '--', repo_url, destination_dir], check=True)
+        print("Cloning successful!")
+    except subprocess.CalledProcessError as e:
+        print("Cloning failed:", e)
+
+# Example usage
+repo_url = "https://huggingface.co/ttmn/SolLlama-mtr"
+destination_dir = "./SolLlama-mtr"
+
+git_clone(repo_url, destination_dir)
+
+import sys
+import os
+import torch
+import numpy as np
+import pandas as pd
+import warnings
+import lightning as L
+torch.set_float32_matmul_precision('high')
+warnings.filterwarnings("ignore", module="pl_bolts")
+
+sys.path.append( '../')
+
+import tokenizer_sl, datamodule_finetune_sl, model_finetune_sl, chemllama_mtr, utils_sl
+import auto_evaluator_sl
+
+torch.manual_seed(1004)
+np.random.seed(1004)
+
+smiles_str = st.text_area('Enter SMILE string')
+file_path = './smiles_str.txt'
+
+# Open the file in write mode ('w') and write the content
+with open(file_path, 'w') as file:
+    file.write(smiles_str)
+
+# smiles_str = "CC02"
+
+###
+# solute_or_solvent = 'solute'
+solute_or_solvent = st.selectbox('Solute or Solvent', ['Solute', 'Solvent'])
+ver_ft = 0 # version control for FT model & evaluation data # Or it will overwrite the models and results
+batch_size_pair = [64, 64] if solute_or_solvent == 'Solute' else [10, 10] # [train, valid(test)] 
+# since 'solute' has very small dataset. So I thinl 10 for train and 10 for valid(test) should be the maximum values.
+lr = 0.0001 
+epochs = 7
+use_freeze = False  # Freeze the model or not # False measn not freezing
+overwrite_level_2 = True
+###
+max_seq_length = 512
+tokenizer = tokenizer_sl.fn_load_tokenizer_llama(
+    max_seq_length=max_seq_length,
+)
+max_length = max_seq_length
+num_workers = 2
+
+# I just reused our previous research code with some modifications.
+dir_main = "."
+name_model_mtr = "ChemLlama_Medium_30m_vloss_val_loss=0.029_ep_epoch=04.ckpt" 
+
+dir_model_mtr = f"{dir_main}/SolLlama-mtr/{name_model_mtr}"
+
+max_seq_length = 512
+
+tokenizer = tokenizer_sl.fn_load_tokenizer_llama(
+    max_seq_length=max_seq_length,
+)
+max_length = max_seq_length
+num_workers = 2
+
+## FT
+
+ver_ft = 0
+dir_model_ft_to_save = f"{dir_main}/SolLlama-mtr"
+# name_model_ft = 'Solvent.pt'
+name_model_ft = f"{solute_or_solvent}.pt"
+
+# Load dataset for finetune
+batch_size_for_train = batch_size_pair[0]
+batch_size_for_valid = batch_size_pair[1]
+
+data_module = datamodule_finetune_sl.CustomFinetuneDataModule(
+    solute_or_solvent=solute_or_solvent,
+    tokenizer=tokenizer,
+    max_seq_length=max_length,
+    batch_size_train=batch_size_for_train,
+    batch_size_valid=batch_size_for_valid,
+    # num_device=int(config.NUM_DEVICE) * config.NUM_WORKERS_MULTIPLIER,
+    num_device=num_workers,
+)
+
+data_module.prepare_data()
+data_module.setup()
+steps_per_epoch = len(data_module.test_dataloader())
+
+# Load model and optimizer for finetune
+learning_rate = lr
+
+
+model_mtr = chemllama_mtr.ChemLlama.load_from_checkpoint(dir_model_mtr)
+
+
+model_ft = model_finetune_sl.CustomFinetuneModel(
+    model_mtr=model_mtr,
+    steps_per_epoch=steps_per_epoch,
+    warmup_epochs=1,
+    max_epochs=epochs,
+    learning_rate=learning_rate,
+    # dataset_dict=dataset_dict,
+    use_freeze=use_freeze,
+)
+
+# 'SolLlama_solute_vloss_val_loss=0.082_ep_epoch=06.pt'
+
+trainer = L.Trainer(
+    default_root_dir=dir_model_ft_to_save,
+    # profiler=profiler,
+    # logger=csv_logger,
+    accelerator='auto',
+    devices='auto',
+    # accelerator='gpu',
+    # devices=[0],
+    min_epochs=1,
+    max_epochs=epochs,
+    precision=32,
+    # callbacks=[checkpoint_callback]
+)
+
+
+# Predict
+local_model_ft = utils_sl.load_model_ft_with(
+    class_model_ft=model_ft, 
+    dir_model_ft=dir_model_ft_to_save,
+    name_model_ft=name_model_ft
+)
+
+result = trainer.predict(local_model_ft, data_module)
+result_pred = list()
+result_label = list()
+for bat in range(len(result)):
+    result_pred.append(result[bat][0].squeeze())
+    result_label.append(result[bat][1])
+
+st.write(result_pred)

.ipynb_checkpoints/chemllama_mtr-checkpoint.py

@@ -0,0 +1,212 @@
+import lightning as L
+import torch
+import torchmetrics
+
+from torch import nn
+from transformers import LlamaModel, LlamaConfig
+from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
+
+class ChemLlama(L.LightningModule):
+    def __init__(
+        self,
+        max_position_embeddings,
+        vocab_size,
+        pad_token_id,
+        bos_token_id,
+        eos_token_id,
+        steps_per_epoch=None, #
+        warmup_epochs=None, #
+        max_epochs=None, #
+        hidden_size=384,
+        intermediate_size=464,
+        num_labels=105,
+        attention_dropout=0.144,
+        num_hidden_layers=3,
+        num_attention_heads=12,
+        learning_rate=0.0001,
+    ):
+        super(ChemLlama, self).__init__()
+        self.save_hyperparameters()
+        
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_labels = num_labels
+        self.vocab_size = vocab_size
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.steps_per_epoch = steps_per_epoch #
+        self.warmup_epochs = warmup_epochs #
+        self.max_epochs = max_epochs #
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.attention_dropout = attention_dropout
+        self.max_position_embeddings = max_position_embeddings
+        self.learning_rate = learning_rate
+
+        self.mae = torchmetrics.MeanAbsoluteError()
+        self.mse = torchmetrics.MeanSquaredError()
+
+        self.config_llama = LlamaConfig(
+            max_position_embeddings=self.max_position_embeddings,
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            intermediate_size=self.intermediate_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            attention_dropout=self.attention_dropout,
+            pad_token_id=self.pad_token_id,
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+        )
+
+        self.loss_fn = nn.L1Loss()
+
+        self.llama = LlamaModel(self.config_llama)
+        self.gelu = nn.GELU()
+        self.score = nn.Linear(self.hidden_size, self.num_labels)
+
+    def forward(self, input_ids, attention_mask, labels=None):
+
+        transformer_outputs = self.llama(
+            input_ids=input_ids, attention_mask=attention_mask
+        )
+
+        hidden_states = transformer_outputs[0]
+        hidden_states = self.gelu(hidden_states)
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config_llama.pad_token_id is None and batch_size != 1:
+            raise ValueError(
+                "Cannot handle batch sizes > 1 if no padding token is defined."
+            )
+        if self.config_llama.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+                sequence_lengths = (
+                    torch.eq(input_ids, self.config_llama.pad_token_id).int().argmax(-1)
+                    - 1
+                )
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+        # raise ValueError(len(sequence_lengths), sequence_lengths)
+
+        pooled_logits = logits[
+            torch.arange(batch_size, device=logits.device), sequence_lengths
+        ]
+        return pooled_logits
+
+    def training_step(self, batch, batch_idx):
+
+        loss, logits, labels = self._common_step(batch=batch, batch_idx=batch_idx)
+
+        # mae = self.mae(logits, labels)
+        # mse = self.mse(logits, labels)
+        self.log_dict(
+            {
+                "train_loss": loss, 
+                # "train_mae": mae, 
+                # "train_mse": mse
+            },
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+            # logger=True,
+        )
+        # on_stecp = True will use lots of computational resources
+
+        # return loss
+        return {"loss": loss, "logits": logits, "labels": labels}
+
+    def train_epoch_end(self, outputs):
+        # avg_loss = torch.stack([x["loss"] for x in outputs]).mean()
+        scores = torch.cat([x["logits"] for x in outputs])
+        labels = torch.cat([x["labels"] for x in outputs])
+        self.log_dict(
+            {
+                "train_mae": self.mae(scores, labels),
+                "train_mse": self.mse(scores, labels)
+            },
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+
+    def validation_step(self, batch, batch_idx):
+
+        loss, logits, labels = self._common_step(batch=batch, batch_idx=batch_idx)
+        # self.log("val_loss", loss)
+        self.log("val_loss", loss, sync_dist=True)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+
+        loss, logits, labels = self._common_step(batch=batch, batch_idx=batch_idx)
+        # self.log("val_loss", loss)
+        self.log("test_loss", loss, sync_dist=True,)
+        return loss
+
+    def _common_step(self, batch, batch_idx):
+
+        logits = self.forward(
+            input_ids=batch["input_ids"].squeeze(),
+            attention_mask=batch["attention_mask"].squeeze(),
+        )
+
+        labels = batch["labels"].squeeze()
+        loss = self.loss_fn(logits, labels)
+
+        # print(f"logits : {logits.shape}")
+        # print(f"labels : {labels.shape}")
+
+        return loss, logits, labels
+
+    # def configure_optimizers(self):  # Schedular here too!
+    #     # since confiture_optimizers and the model are included in the same class.. self.parameters()
+    #     return torch.optim.AdamW(
+    #         params=self.parameters(),
+    #         lr=self.learning_rate,
+    #         betas=(0.9, 0.999),
+    #         weight_decay=0.01,
+    #     )
+
+    # # The below is for warm-up scheduler
+    # https://lightning.ai/forums/t/how-to-use-warmup-lr-cosineannealinglr-in-lightning/1980
+    # https://github.com/Lightning-AI/pytorch-lightning/issues/328
+    def configure_optimizers(self):  # Schedular here too!
+        # since confiture_optimizers and the model are included in the same class.. self.parameters()
+        optimizer = torch.optim.AdamW(
+            params=self.parameters(),
+            lr=self.learning_rate,
+            betas=(0.9, 0.999),
+            weight_decay=0.01,
+        )
+
+        # "warmup_epochs //4 only not max_epochs" will work
+        scheduler = LinearWarmupCosineAnnealingLR(
+            optimizer, 
+            warmup_epochs=self.warmup_epochs*self.steps_per_epoch // torch.cuda.device_count(), # // num_device in this case
+            max_epochs=self.max_epochs*self.steps_per_epoch // torch.cuda.device_count(),
+        )
+
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "interval": "step",
+                "frequency": 1,
+                "reduce_on_plateau": False,
+                "monitor": "val_loss",
+            }
+        }

app.py

@@ -33,6 +33,9 @@ sys.path.append( '../')
 import tokenizer_sl, datamodule_finetune_sl, model_finetune_sl, chemllama_mtr, utils_sl
 import auto_evaluator_sl
 
+from torch.utils.data import Dataset, DataLoader
+from transformers import DataCollatorWithPadding
+
 torch.manual_seed(1004)
 np.random.seed(1004)
 
@@ -48,26 +51,143 @@ with open(file_path, 'w') as file:
 ###
 # solute_or_solvent = 'solute'
 solute_or_solvent = st.selectbox('Solute or Solvent', ['Solute', 'Solvent'])
-ver_ft = 0 # version control for FT model & evaluation data # Or it will overwrite the models and results
-batch_size_pair = [64, 64] if solute_or_solvent == 'Solute' else [10, 10] # [train, valid(test)] 
-# since 'solute' has very small dataset. So I thinl 10 for train and 10 for valid(test) should be the maximum values.
-lr = 0.0001 
-epochs = 7
-use_freeze = False  # Freeze the model or not # False measn not freezing
-overwrite_level_2 = True
-###
-max_seq_length = 512
-tokenizer = tokenizer_sl.fn_load_tokenizer_llama(
-    max_seq_length=max_seq_length,
-)
-max_length = max_seq_length
-num_workers = 2
+
+
+class ChemLlama(nn.Module):
+    def __init__(
+        self,
+        max_position_embeddings=512,
+        vocab_size=591,
+        pad_token_id=0,
+        bos_token_id=12,
+        eos_token_id=13,
+        hidden_size=768,
+        intermediate_size=768,
+        num_labels=105,
+        attention_dropout=0.144,
+        num_hidden_layers=7,
+        num_attention_heads=8,
+        learning_rate=0.0001,
+    ):
+        super(ChemLlama, self).__init__()
+        
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_labels = num_labels
+        self.vocab_size = vocab_size
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.attention_dropout = attention_dropout
+        self.max_position_embeddings = max_position_embeddings
+
+        self.mae = torchmetrics.MeanAbsoluteError()
+        self.mse = torchmetrics.MeanSquaredError()
+
+        self.config_llama = LlamaConfig(
+            max_position_embeddings=self.max_position_embeddings,
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            intermediate_size=self.intermediate_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            attention_dropout=self.attention_dropout,
+            pad_token_id=self.pad_token_id,
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+        )
+
+        self.loss_fn = nn.L1Loss()
+
+        self.llama = LlamaModel(self.config_llama)
+        self.gelu = nn.GELU()
+        self.score = nn.Linear(self.hidden_size, self.num_labels)
+
+    def forward(self, input_ids, attention_mask, labels=None):
+
+        transformer_outputs = self.llama(
+            input_ids=input_ids, attention_mask=attention_mask
+        )
+
+        hidden_states = transformer_outputs[0]
+        hidden_states = self.gelu(hidden_states)
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config_llama.pad_token_id is None and batch_size != 1:
+            raise ValueError(
+                "Cannot handle batch sizes > 1 if no padding token is defined."
+            )
+        if self.config_llama.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+                sequence_lengths = (
+                    torch.eq(input_ids, self.config_llama.pad_token_id).int().argmax(-1)
+                    - 1
+                )
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+        # raise ValueError(len(sequence_lengths), sequence_lengths)
+
+        pooled_logits = logits[
+            torch.arange(batch_size, device=logits.device), sequence_lengths
+        ]
+        return pooled_logits
+
+
+chemllama_mtr = ChemLlama()
+
+class ChemLlama_FT(nn.Module):
+    def __init__(
+        self,
+        model_mtr,
+        linear_param:int=64,
+        use_freeze:bool=True,
+        *args, **kwargs
+    ):
+        super(CustomFinetuneModel, self).__init__()
+        # self.save_hyperparameters()
+
+        self.model_mtr = model_mtr
+        if use_freeze:
+            self.model_mtr.freeze()
+            # for name, param in model_mtr.named_parameters():
+            #     param.requires_grad = False
+            #     print(name, param.requires_grad)
+        
+        self.gelu = nn.GELU()
+        self.linear1 = nn.Linear(self.model_mtr.num_labels, linear_param)
+        self.linear2 = nn.Linear(linear_param, linear_param)
+        self.regression = nn.Linear(linear_param, 5)
+
+        self.loss_fn = nn.L1Loss()
+
+    def forward(self, input_ids, attention_mask, labels=None):
+        x = self.model_mtr(input_ids=input_ids, attention_mask=attention_mask)
+        x = self.gelu(x)
+        x = self.linear1(x)
+        x = self.gelu(x)
+        x = self.linear2(x)
+        x = self.gelu(x)
+        x = self.regression(x)
+        
+        return x
+        
+chemllama_ft = ChemLlama_FT(model_mtr=chemllama_mtr)
+
 
 # I just reused our previous research code with some modifications.
 dir_main = "."
-name_model_mtr = "ChemLlama_Medium_30m_vloss_val_loss=0.029_ep_epoch=04.ckpt" 
-
-dir_model_mtr = f"{dir_main}/SolLlama-mtr/{name_model_mtr}"
 
 max_seq_length = 512
 
@@ -79,75 +199,95 @@ num_workers = 2
 
 ## FT
 
-ver_ft = 0
 dir_model_ft_to_save = f"{dir_main}/SolLlama-mtr"
 # name_model_ft = 'Solvent.pt'
 name_model_ft = f"{solute_or_solvent}.pt"
 
-# Load dataset for finetune
-batch_size_for_train = batch_size_pair[0]
-batch_size_for_valid = batch_size_pair[1]
-
-data_module = datamodule_finetune_sl.CustomFinetuneDataModule(
-    solute_or_solvent=solute_or_solvent,
-    tokenizer=tokenizer,
-    max_seq_length=max_length,
-    batch_size_train=batch_size_for_train,
-    batch_size_valid=batch_size_for_valid,
-    # num_device=int(config.NUM_DEVICE) * config.NUM_WORKERS_MULTIPLIER,
-    num_device=num_workers,
-)
-
-data_module.prepare_data()
-data_module.setup()
-steps_per_epoch = len(data_module.test_dataloader())
-
-# Load model and optimizer for finetune
-learning_rate = lr
-
-
-model_mtr = chemllama_mtr.ChemLlama.load_from_checkpoint(dir_model_mtr)
-
-
-model_ft = model_finetune_sl.CustomFinetuneModel(
-    model_mtr=model_mtr,
-    steps_per_epoch=steps_per_epoch,
-    warmup_epochs=1,
-    max_epochs=epochs,
-    learning_rate=learning_rate,
-    # dataset_dict=dataset_dict,
-    use_freeze=use_freeze,
-)
-
-# 'SolLlama_solute_vloss_val_loss=0.082_ep_epoch=06.pt'
-
-trainer = L.Trainer(
-    default_root_dir=dir_model_ft_to_save,
-    # profiler=profiler,
-    # logger=csv_logger,
-    accelerator='auto',
-    devices='auto',
-    # accelerator='gpu',
-    # devices=[0],
-    min_epochs=1,
-    max_epochs=epochs,
-    precision=32,
-    # callbacks=[checkpoint_callback]
-)
-
-
+# # Load dataset for finetune
+# batch_size_for_train = batch_size_pair[0]
+# batch_size_for_valid = batch_size_pair[1]
+
+# data_module = datamodule_finetune_sl.CustomFinetuneDataModule(
+#     solute_or_solvent=solute_or_solvent,
+#     tokenizer=tokenizer,
+#     max_seq_length=max_length,
+#     batch_size_train=batch_size_for_train,
+#     batch_size_valid=batch_size_for_valid,
+#     # num_device=int(config.NUM_DEVICE) * config.NUM_WORKERS_MULTIPLIER,
+#     num_device=num_workers,
+# )
+
+# data_module.prepare_data()
+# data_module.setup()
+# steps_per_epoch = len(data_module.test_dataloader())
+
+# # Load model and optimizer for finetune
+# learning_rate = lr
+
+
+# model_mtr = chemllama_mtr.ChemLlama.load_from_checkpoint(dir_model_mtr)
+
+
+# model_ft = model_finetune_sl.CustomFinetuneModel(
+#     model_mtr=model_mtr,
+#     steps_per_epoch=steps_per_epoch,
+#     warmup_epochs=1,
+#     max_epochs=epochs,
+#     learning_rate=learning_rate,
+#     # dataset_dict=dataset_dict,
+#     use_freeze=use_freeze,
+# )
+
+# # 'SolLlama_solute_vloss_val_loss=0.082_ep_epoch=06.pt'
+
+# trainer = L.Trainer(
+#     default_root_dir=dir_model_ft_to_save,
+#     # profiler=profiler,
+#     # logger=csv_logger,
+#     accelerator='auto',
+#     devices='auto',
+#     # accelerator='gpu',
+#     # devices=[0],
+#     min_epochs=1,
+#     max_epochs=epochs,
+#     precision=32,
+#     # callbacks=[checkpoint_callback]
+# )
+
+device = 'cpu'
 # Predict
 local_model_ft = utils_sl.load_model_ft_with(
-    class_model_ft=model_ft, 
+    class_model_ft=chemllama_ft, 
     dir_model_ft=dir_model_ft_to_save,
     name_model_ft=name_model_ft
+).to(device)
+
+# result = trainer.predict(local_model_ft, data_module)
+# result_pred = list()
+# result_label = list()
+# for bat in range(len(result)):
+#     result_pred.append(result[bat][0].squeeze())
+#     result_label.append(result[bat][1])
+
+with open('./smiles_str.txt', 'r') as file:
+    smiles_str = file.readline()
+    
+dataset_test = datamodule_finetune_sl.CustomLlamaDatasetAbraham(
+    df=pd.DataFrame([smiles_str]),
+    tokenizer=tokenizer,
+    max_seq_length=max_length
 )
 
-result = trainer.predict(local_model_ft, data_module)
-result_pred = list()
-result_label = list()
-for bat in range(len(result)):
-    result_pred.append(result[bat][0].squeeze())
-    result_label.append(result[bat][1])
+dataloader_test = DataLoader(dataset_test, shuffle=False, collate_fn=DataCollatorWithPadding(tokenizer))
+
+list_predictions = []
+local_model_ft.eval()
+with torch.inference_mode():
+    for i, v_batch in enumerate(dataloader_test):
+        v_input_ids = v_batch['input_ids'].to(device)
+        v_attention_mask = v_batch['attention_mask'].to(device)
+        # v_y_labels = v_batch['labels'].to(device)
+        v_y_logits = local_model_ft(input_ids=v_input_ids, attention_mask=v_attention_mask)
+        list_predictions.append(v_y_logits[0][0].tolist())
 
-st.write(result_pred)
+st.write(list_predictions)

app_old.py

@@ -0,0 +1,153 @@
+
+import streamlit as st
+# from git import Repo
+
+# Repo.clone_from('https://huggingface.co/ttmn/SolLlama-mtr', './SolLlama-mtr')
+import subprocess
+
+def git_clone(repo_url, destination_dir):
+    try:
+        subprocess.run(['git', 'clone', '-v', '--', repo_url, destination_dir], check=True)
+        print("Cloning successful!")
+    except subprocess.CalledProcessError as e:
+        print("Cloning failed:", e)
+
+# Example usage
+repo_url = "https://huggingface.co/ttmn/SolLlama-mtr"
+destination_dir = "./SolLlama-mtr"
+
+git_clone(repo_url, destination_dir)
+
+import sys
+import os
+import torch
+import numpy as np
+import pandas as pd
+import warnings
+import lightning as L
+torch.set_float32_matmul_precision('high')
+warnings.filterwarnings("ignore", module="pl_bolts")
+
+sys.path.append( '../')
+
+import tokenizer_sl, datamodule_finetune_sl, model_finetune_sl, chemllama_mtr, utils_sl
+import auto_evaluator_sl
+
+torch.manual_seed(1004)
+np.random.seed(1004)
+
+smiles_str = st.text_area('Enter SMILE string')
+file_path = './smiles_str.txt'
+
+# Open the file in write mode ('w') and write the content
+with open(file_path, 'w') as file:
+    file.write(smiles_str)
+
+# smiles_str = "CC02"
+
+###
+# solute_or_solvent = 'solute'
+solute_or_solvent = st.selectbox('Solute or Solvent', ['Solute', 'Solvent'])
+ver_ft = 0 # version control for FT model & evaluation data # Or it will overwrite the models and results
+batch_size_pair = [64, 64] if solute_or_solvent == 'Solute' else [10, 10] # [train, valid(test)] 
+# since 'solute' has very small dataset. So I thinl 10 for train and 10 for valid(test) should be the maximum values.
+lr = 0.0001 
+epochs = 7
+use_freeze = False  # Freeze the model or not # False measn not freezing
+overwrite_level_2 = True
+###
+max_seq_length = 512
+tokenizer = tokenizer_sl.fn_load_tokenizer_llama(
+    max_seq_length=max_seq_length,
+)
+max_length = max_seq_length
+num_workers = 2
+
+# I just reused our previous research code with some modifications.
+dir_main = "."
+name_model_mtr = "ChemLlama_Medium_30m_vloss_val_loss=0.029_ep_epoch=04.ckpt" 
+
+dir_model_mtr = f"{dir_main}/SolLlama-mtr/{name_model_mtr}"
+
+max_seq_length = 512
+
+tokenizer = tokenizer_sl.fn_load_tokenizer_llama(
+    max_seq_length=max_seq_length,
+)
+max_length = max_seq_length
+num_workers = 2
+
+## FT
+
+ver_ft = 0
+dir_model_ft_to_save = f"{dir_main}/SolLlama-mtr"
+# name_model_ft = 'Solvent.pt'
+name_model_ft = f"{solute_or_solvent}.pt"
+
+# Load dataset for finetune
+batch_size_for_train = batch_size_pair[0]
+batch_size_for_valid = batch_size_pair[1]
+
+data_module = datamodule_finetune_sl.CustomFinetuneDataModule(
+    solute_or_solvent=solute_or_solvent,
+    tokenizer=tokenizer,
+    max_seq_length=max_length,
+    batch_size_train=batch_size_for_train,
+    batch_size_valid=batch_size_for_valid,
+    # num_device=int(config.NUM_DEVICE) * config.NUM_WORKERS_MULTIPLIER,
+    num_device=num_workers,
+)
+
+data_module.prepare_data()
+data_module.setup()
+steps_per_epoch = len(data_module.test_dataloader())
+
+# Load model and optimizer for finetune
+learning_rate = lr
+
+
+model_mtr = chemllama_mtr.ChemLlama.load_from_checkpoint(dir_model_mtr)
+
+
+model_ft = model_finetune_sl.CustomFinetuneModel(
+    model_mtr=model_mtr,
+    steps_per_epoch=steps_per_epoch,
+    warmup_epochs=1,
+    max_epochs=epochs,
+    learning_rate=learning_rate,
+    # dataset_dict=dataset_dict,
+    use_freeze=use_freeze,
+)
+
+# 'SolLlama_solute_vloss_val_loss=0.082_ep_epoch=06.pt'
+
+trainer = L.Trainer(
+    default_root_dir=dir_model_ft_to_save,
+    # profiler=profiler,
+    # logger=csv_logger,
+    accelerator='auto',
+    devices='auto',
+    # accelerator='gpu',
+    # devices=[0],
+    min_epochs=1,
+    max_epochs=epochs,
+    precision=32,
+    # callbacks=[checkpoint_callback]
+)
+
+
+# Predict
+local_model_ft = utils_sl.load_model_ft_with(
+    class_model_ft=model_ft, 
+    dir_model_ft=dir_model_ft_to_save,
+    name_model_ft=name_model_ft
+)
+
+result = trainer.predict(local_model_ft, data_module)
+result_pred = list()
+result_label = list()
+for bat in range(len(result)):
+    result_pred.append(result[bat][0].squeeze())
+    result_label.append(result[bat][1])
+
+st.write(result_pred)