Datawhale AI夏令营（笔记②）

#Datawhale AI夏令营# #AI# #夏令营#

        作为一个新手小白，这是我第一次接触AI与化学结合的赛题，以及第一次接触机器学习与深度学习。

第一步：学习课程内容

一.机器学习

        机器学习按照目标可以分为分类任务（classification）和回归（regression）任务两大类。所谓分类任务，就是模型预测的结果是离散的值，例如类别；那么，回归任务中，模型预测的结果就是连续的值，例如房价等等。以下是一些概念的基本科普：

        1. 分类任务（Classification）

分类任务的目标是将数据点分配到预定义的类别中。模型的输出是离散的标签值。常见的分类算法包括：

• 逻辑回归（Logistic Regression）: 虽名为回归算法，但用于分类问题，通过计算事件发生的概率来进行分类。
• 支持向量机（SVM, Support Vector Machine）: 通过找到一个最优超平面来区分不同的类别。
• 决策树（Decision Tree）: 通过树状结构进行决策，逐步分割数据来达到分类的目的。
• 随机森林（Random Forest）: 由多棵决策树组成，通过对每棵树的预测结果进行投票来得到最终的分类结果。
• K-近邻（K-Nearest Neighbors, KNN）: 基于距离度量对新数据点进行分类，通过查看其最近邻居的类别来做出预测。
• 神经网络（Neural Networks）: 通过多层网络学习复杂的特征来进行分类，特别适合处理大规模数据集。

        2. 回归任务（Regression）

回归任务的目标是预测一个连续的数值。模型的输出是一个连续值。常见的回归算法包括：

• 线性回归（Linear Regression）: 假设输入特征与输出之间存在线性关系，通过最小化误差来找到最佳拟合直线。
• 岭回归（Ridge Regression）: 在线性回归的基础上加入L2正则化项，以减少模型的过拟合。
• Lasso回归（Lasso Regression）: 在线性回归中加入L1正则化项，有助于特征选择。
• 决策树回归（Decision Tree Regression）: 使用决策树来进行回归预测，通过分裂数据点来减少预测误差。
• 随机森林回归（Random Forest Regression）: 多棵决策树的集合，通过平均每棵树的预测结果来提高回归性能。
• 梯度提升回归（Gradient Boosting Regression）: 通过逐步训练模型并修正前一模型的误差来提高预测精度。
• 神经网络回归（Neural Network Regression）: 通过深度学习模型来学习复杂的非线性关系进行回归预测。

        3. 特征工程（Feature Engineering）

特征工程是机器学习中的一个重要步骤，涉及到从原始数据中提取特征，形成模型可以使用的输入数据。常见的特征工程方法包括：

• 数据清洗（Data Cleaning）: 处理缺失值、异常值等数据问题。
• 特征选择（Feature Selection）: 选择对目标变量有重要影响的特征，去除冗余或无关特征。
• 特征转换（Feature Transformation）: 将原始特征转换为更有用的特征，比如通过标准化、归一化、对数变换等方法。
• 特征提取（Feature Extraction）: 从原始数据中提取新的特征，比如通过主成分分析（PCA）或特征构造（如多项式特征）等方法。

        4. 决策树（Decision Tree）

决策树是一种监督学习算法，可以用于分类或回归任务。决策树通过树状结构来决策，每个节点表示对某个特征的判断，每个分支表示判断结果，每个叶子节点表示最终的预测结果。决策树的优点是易于解释，但容易过拟合。

• 信息熵（Entropy）: 用于衡量分类任务中的不确定性，通过选择信息增益最大的特征来分裂节点。
• 均方误差（Mean Squared Error, MSE）: 用于回归任务，通过最小化预测值和真实值之间的均方差来决定节点的分裂。

•  传统的机器学习需要需要经历特征工程这一步骤，即将原始数据转化为向量形式。然后通过SVM、Random Forest等算法学习数据的规律。这些方法在处理简单的任务时是比较有效的。

          图3 决策树 （左）分类型决策树，（右）回归型决策树

          划分每个节点的目标是让该节点中的值尽可能相同。在分类任务中，常见的就是信息熵衡量；在回归任务中，可以使用均方误差、绝对误差等进行衡量。

       

        5. 随机森林（Random Forest）

随机森林是一种集成学习方法，通过构建多棵决策树并将其预测结果进行汇总来提高模型的准确性和鲁棒性。每棵决策树在训练时使用的数据集是从原始数据集中随机抽取的子集，且每个节点的特征选择也是随机的。随机森林通过对每棵树的预测结果进行投票（分类）或平均（回归）来得到最终的预测结果。

• 优点: 可以有效处理高维数据和特征之间的非线性关系，具有较高的准确性和较低的过拟合风险。
• 缺点: 计算成本较高，模型解释性较差

        在本次竞赛中，我们需要预测的目标是反应的产率，是0-1之间的一个连续的数值，所以是一个回归任务。（注：离散值通过一些处理可以近似认为是连续值，所以不要被连续值和离散值限制了自己的思维）。

        传统的机器学习需要需要经历特征工程这一步骤，即将原始数据转化为向量形式。然后通过SVM、Random Forest等算法学习数据的规律。这些方法在处理简单的任务时是比较有效的。

         图4 随机森林

        将多个决策树结合在一起，训练每个决策树的数据集都是随机有放回地从原数据中选出。预测的时候，输入会通过每个决策树进行预测，然后考虑每个树地输出结果，得到最终的预测值。

二.深度学习

        深度学习可以归为机器学习的一个子集，主要通过神经网络学习数据的特征和分布。深度学习的一个重要进化是不再需要繁琐的特征工程，让神经网络自己从里面学习特征。

        SMILES是一种以ASCII组成的序列，可以被理解为一种“化学语言”。既然是一种语言，那么很自然地想到了可以使用NLP中的方法对SMILES进行建模。

        使用RNN对SMILES建模是早期的一个主要方法。RNN（Recurrent Neural Network）是处理序列数据的一把好手。RNN的网络每层除了会有自己的输出以外，还会输出一个隐向量到下一层。

1.RNN的架构与公式

RNN的基本构建块包括：

• 隐状态（Hidden State, hnh_nhn​）: 用于存储前面序列的信息，并传递到当前时间步。
• 输入门（Input Gate）: 控制当前输入对隐状态的影响。
• 遗忘门（Forget Gate）: 控制上一个隐状态的遗忘程度。
• 输出门（Output Gate）: 控制隐状态对最终输出的影响。

公式：

1. 隐状态更新:

2. 输出计算:

       

  图5 RNN的架构示意图

        其中，每一层相当于做了一次线性变换：

   

        通过隐向量的不断传递，序列后面的部分就通过“阅读”隐向量，获取前面序列的信息，从而提升学习能力。

       2.RNN的作用

1.  特征嵌入与增强：
   • RNN可以被用作一种特征提取器或嵌入层，用于将反应底物、添加剂、溶剂等的信息转换成更高级别的表示。这些表示随后可以被用作其他机器学习或深度学习模型（如全连接神经网络、卷积神经网络等）的输入。
2. 结合文本信息：
   • 如果反应数据集中包含有关反应条件、底物结构或催化剂的文本描述，RNN（特别是结合注意力机制的LSTM或GRU）可以用于提取这些文本中的关键信息，并将其转换为数值特征，这些特征可以用于预测产率。

     

       3.RNN的缺点 

        如果序列太长，那么两个相距比较远的字符之间的联系需要通过多个隐藏向量。这就像人和人之间传话一样，传递的人多了，很容易导致信息的损失或者扭曲。因此，它对长序列的记忆能力较弱。

        同时，RNN需要一层一层地传递，所以并行能力差，同时也比较容易出现梯度消失或梯度爆炸问题。（据说Transfer会好很多，明天再学习新内容）

第二步：理解实践代码部分

一.定义RNN模型

# 定义RNN模型
class RNNModel(nn.Module):
    def __init__(self, num_embed, input_size, hidden_size, output_size, num_layers, dropout, device):
        super(RNNModel, self).__init__()
        self.embed = nn.Embedding(num_embed, input_size)
        self.rnn = nn.RNN(input_size, hidden_size, num_layers=num_layers, 
                          batch_first=True, dropout=dropout, bidirectional=True)
        self.fc = nn.Sequential(nn.Linear(2 * num_layers * hidden_size, output_size),
                                nn.Sigmoid(),
                                nn.Linear(output_size, 1),
                                nn.Sigmoid())

    def forward(self, x):
        # x : [bs, seq_len]
        x = self.embed(x)
        # x : [bs, seq_len, input_size]
        _, hn = self.rnn(x) # hn : [2*num_layers, bs, h_dim]
        hn = hn.transpose(0,1)
        z = hn.reshape(hn.shape[0], -1) # z shape: [bs, 2*num_layers*h_dim]
        output = self.fc(z).squeeze(-1) # output shape: [bs, 1]
        return output

1.RNN模型理解

初始化函数：__init__

       在初始化函数中定义了这个RNN模型需要的所有组件。

• 嵌入层 (self.embed): 这个层负责将输入的整数（比如单词的索引）转换成固定大小的密集向量（embedding vectors）。num_embed是嵌入层的大小（即词汇表的大小），input_size是每个嵌入向量的维度。

• RNN层 (self.rnn): 这是RNN模型的核心层。它接受嵌入层的输出作为输入，并沿着时间步（或序列中的元素）进行迭代处理。input_size是输入向量的维度，hidden_size是隐藏层的维度（即RNN单元在每个时间步的内部状态的大小），num_layers是RNN的层数（可以堆叠多个RNN层来增加网络的深度），dropout用于减少过拟合，bidirectional=True表示RNN是双向的，即它会同时处理序列的正向和反向，这有助于捕获更多的上下文信息。

• 全连接层 (self.fc): 这个层将RNN层的输出转换成最终的预测结果。由于RNN是双向的，所以它的输出维度是2 * num_layers * hidden_size（两层输出，每层都有num_layers * hidden_size个单元）。这里使用了两层线性层（中间加了一个Sigmoid激活函数），最后输出层是一个线性层后跟Sigmoid激活函数，用于将输出压缩到0和1之间（这通常用于二分类问题）。

前向传播函数：forward

      这个函数定义了数据通过网络的前向传播路径。

1. 嵌入输入：首先，输入的整数序列（x）通过嵌入层转换成向量序列。

2. RNN处理：然后，这些向量序列被送入RNN层进行处理。这里只关心RNN层的最后一个输出状态（hn），因为它包含了整个序列的信息。注意，由于RNN是双向的，所以hn的形状是[2*num_layers, bs, h_dim]，其中bs是批次大小，h_dim是隐藏层的维度。

3. 调整RNN输出：接下来是调整hn的形状以便将其送入全连接层。首先交换hn的第一和第二个维度（hn.transpose(0,1)），使其形状变为[bs, 2*num_layers, h_dim]。然后，我们将这个三维张量展平成一个二维张量（z），其形状为[bs, 2*num_layers*h_dim]。

4. 全连接层处理：最后代码将z送入全连接层，得到最终的预测结果。由于输出层后面跟了一个squeeze(-1)，所以最终输出的形状是[bs, 1]，即每个样本一个预测值。

二.数据处理

# import matplotlib.pyplot as plt
## 数据处理部分
# tokenizer，鉴于SMILES的特性，这里需要自己定义tokenizer和vocab
# 这里直接将smiles str按字符拆分，并替换为词汇表中的序号
class Smiles_tokenizer():
    def __init__(self, pad_token, regex, vocab_file, max_length):
        self.pad_token = pad_token
        self.regex = regex
        self.vocab_file = vocab_file
        self.max_length = max_length

        with open(self.vocab_file, "r") as f:
            lines = f.readlines()
        lines = [line.strip("\n") for line in lines]
        vocab_dic = {}
        for index, token in enumerate(lines):
            vocab_dic[token] = index
        self.vocab_dic = vocab_dic

    def _regex_match(self, smiles):
        regex_string = r"(" + self.regex + r"|"
        regex_string += r".)"
        prog = re.compile(regex_string)

        tokenised = []
        for smi in smiles:
            tokens = prog.findall(smi)
            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length]
            tokenised.append(tokens) # 返回一个所有的字符串列表
        return tokenised
    
    def tokenize(self, smiles):
        tokens = self._regex_match(smiles)
        # 添加上表示开始和结束的token：<cls>, <end>
        tokens = [["<CLS>"] + token + ["<SEP>"] for token in tokens]
        tokens = self._pad_seqs(tokens, self.pad_token)
        token_idx = self._pad_token_to_idx(tokens)
        return tokens, token_idx

    def _pad_seqs(self, seqs, pad_token):
        pad_length = max([len(seq) for seq in seqs])
        padded = [seq + ([pad_token] * (pad_length - len(seq))) for seq in seqs]
        return padded

    def _pad_token_to_idx(self, tokens):
        idx_list = []
        for token in tokens:
            tokens_idx = []
            for i in token:
                if i in self.vocab_dic.keys():
                    tokens_idx.append(self.vocab_dic[i])
                else:
                    self.vocab_dic[i] = max(self.vocab_dic.values()) + 1
                    tokens_idx.append(self.vocab_dic[i])
            idx_list.append(tokens_idx)
        
        return idx_list

# 读数据并处理
def read_data(file_path, train=True):
    df = pd.read_csv(file_path)
    reactant1 = df["Reactant1"].tolist()
    reactant2 = df["Reactant2"].tolist()
    product = df["Product"].tolist()
    additive = df["Additive"].tolist()
    solvent = df["Solvent"].tolist()
    if train:
        react_yield = df["Yield"].tolist()
    else:
        react_yield = [0 for i in range(len(reactant1))]
    
    # 将reactant拼到一起，之间用.分开。product也拼到一起，用>分开
    input_data_list = []
    for react1, react2, prod, addi, sol in zip(reactant1, reactant2, product, additive, solvent):
        input_info = ".".join([react1, react2])
        input_info = ">".join([input_info, prod])
        input_data_list.append(input_info)
    output = [(react, y) for react, y in zip(input_data_list, react_yield)]

    # 下面的代码将reactant\additive\solvent拼到一起，之间用.分开。product也拼到一起，用>分开
    '''
    input_data_list = []
    for react1, react2, prod, addi, sol in zip(reactant1, reactant2, product, additive, solvent):
        input_info = ".".join([react1, react2, addi, sol])
        input_info = ">".join([input_info, prod])
        input_data_list.append(input_info)
    output = [(react, y) for react, y in zip(input_data_list, react_yield)]
    '''

    # # 统计seq length，序列的长度是一个重要的参考，可以使用下面的代码统计查看以下序列长度的分布
    # seq_length = [len(i[0]) for i in output]
    # seq_length_400 = [len(i[0]) for i in output if len(i[0])>200]
    # print(len(seq_length_400) / len(seq_length))
    # seq_length.sort(reverse=True)
    # plt.plot(range(len(seq_length)), seq_length)
    # plt.title("templates frequence")
    # plt.show()
    return output

class ReactionDataset(Dataset):
    def __init__(self, data: List[Tuple[List[str], float]]):
        self.data = data
        
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]
    
def collate_fn(batch):
    REGEX = r"\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9]"
    tokenizer = Smiles_tokenizer("<PAD>", REGEX, "../vocab_full.txt", max_length=300)
    smi_list = []
    yield_list = []
    for i in batch:
        smi_list.append(i[0])
        yield_list.append(i[1])
    tokenizer_batch = torch.tensor(tokenizer.tokenize(smi_list)[1])
    yield_list = torch.tensor(yield_list)
    return tokenizer_batch, yield_list

（知道是干嘛的，但是对写代码过程很懵逼，简单梳理一下）

Smiles_tokenizer 类

1. 初始化 (__init__)

   • pad_token: 用于填充序列的特殊符号。
   • regex: 用于分割SMILES字符串的正则表达式。
   • vocab_file: 词汇表文件路径。
   • max_length: 序列的最大长度。
   • 初始化时会读取词汇表文件并创建词汇字典。

2. 正则匹配 (_regex_match)

   • 使用正则表达式将SMILES字符串分割成tokens。
   • 对每个SMILES字符串进行分割，并且截取到最大长度。

3. Token化 (tokenize)

   • 调用_regex_match方法分割SMILES字符串。
   • 为每个token序列添加开始 (<CLS>) 和结束 (<SEP>) 标记。
   • 使用 _pad_seqs 方法对序列进行填充。
   • 使用 _pad_token_to_idx 方法将tokens转换为对应的索引。

4. 序列填充 (_pad_seqs)

   • 根据最长的序列长度对所有序列进行填充，使其长度相同。

5. tokens转换为索引 (_pad_token_to_idx)

   • 将tokens转换为词汇表中的索引，如果词汇表中没有该token，则为其分配一个新的索引。

数据读取和处理 (read_data)

• 从CSV文件中读取数据，包括反应物、产物、添加剂、溶剂和产率。
• 将反应物和产物拼接成一个字符串，并形成输入数据列表。
• 将输入数据和产率打包成元组列表，作为输出。

数据集类 (ReactionDataset)

• 继承自 torch.utils.data.Dataset，用于创建自定义数据集。
• 实现了 __len__ 和 __getitem__ 方法，分别用于获取数据集长度和获取数据项。

批处理函数 (collate_fn)

• 用于将数据集中的数据打包成批次。
• 定义了用于分割SMILES字符串的正则表达式。
• 实例化 Smiles_tokenizer。
• 对批次中的SMILES字符串进行token化，并转换为索引。
• 将产率转换为张量。

三.训练数据

• 设置超参数。
• 读取并准备数据。
• 创建数据加载器。
• 定义模型、优化器和损失函数。
• 进行训练，打印每轮的损失并保存最佳模型。
• 计算并打印训练总时间。

代码如下：

def train():
    ## super param
    N = 10  #int / int(len(dataset) * 1)  # 或者你可以设置为数据集大小的一定比例，如 int(len(dataset) * 0.1)
    NUM_EMBED = 294 # nn.Embedding()
    INPUT_SIZE = 300 # src length
    HIDDEN_SIZE = 512
    OUTPUT_SIZE = 512
    NUM_LAYERS = 10
    DROPOUT = 0.2
    CLIP = 1 # CLIP value
    N_EPOCHS = 100
    LR = 0.0001
    
    start_time = time.time()  # 开始计时
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # device = 'cpu'
    data = read_data("../dataset/round1_train_data.csv")
    dataset = ReactionDataset(data)
    subset_indices = list(range(N))
    subset_dataset = Subset(dataset, subset_indices)
    train_loader = DataLoader(dataset, batch_size=128, shuffle=True, collate_fn=collate_fn)

    model = RNNModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)
    model.train()
    
    optimizer = optim.Adam(model.parameters(), lr=LR)
    # criterion = nn.MSELoss() # MSE
    criterion = nn.L1Loss() # MAE

    best_loss = 10
    for epoch in range(N_EPOCHS):
        epoch_loss = 0
        for i, (src, y) in enumerate(train_loader):
            src, y = src.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(src)
            loss = criterion(output, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP)
            optimizer.step()
            epoch_loss += loss.item()
            loss_in_a_epoch = epoch_loss / len(train_loader)
        print(f'Epoch: {epoch+1:02} | Train Loss: {loss_in_a_epoch:.3f}')
        if loss_in_a_epoch < best_loss:
            # 在训练循环结束后保存模型
            torch.save(model.state_dict(), '../model/RNN.pth')
    end_time = time.time()  # 结束计时
    # 计算并打印运行时间
    elapsed_time_minute = (end_time - start_time)/60
    print(f"Total running time: {elapsed_time_minute:.2f} minutes")

if __name__ == '__main__':
    train()

超参数这一块我进行了一些参数上的优化：

import time
import torch
import torch.optim as optim
from torch.utils.data import DataLoader, Subset

def train():
    ## 超参数设置
    N = 10
    NUM_EMBED = 294
    INPUT_SIZE = 300
    HIDDEN_SIZE = 256  # 调整为较小的值以减少训练时间
    OUTPUT_SIZE = 256  # 调整为较小的值以减少训练时间
    NUM_LAYERS = 3     # 调整为较小的值以减少训练时间
    DROPOUT = 0.3      # 增大Dropout率来防止过拟合
    CLIP = 0.5         # 调整梯度裁剪的阈值
    N_EPOCHS = 50      # 减少训练轮数
    LR = 0.001         # 增大学习率

    start_time = time.time()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    data = read_data("../dataset/round1_train_data.csv")
    dataset = ReactionDataset(data)
    subset_indices = list(range(N))
    subset_dataset = Subset(dataset, subset_indices)
    train_loader = DataLoader(dataset, batch_size=64, shuffle=True, collate_fn=collate_fn)  # 减小批次大小

    model = RNNModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)
    model.train()
    
    optimizer = optim.Adam(model.parameters(), lr=LR)
    criterion = nn.L1Loss()

    best_loss = 10
    for epoch in range(N_EPOCHS):
        epoch_loss = 0
        for i, (src, y) in enumerate(train_loader):
            src, y = src.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(src)
            loss = criterion(output, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP)
            optimizer.step()
            epoch_loss += loss.item()
        loss_in_a_epoch = epoch_loss / len(train_loader)
        print(f'Epoch: {epoch+1:02} | Train Loss: {loss_in_a_epoch:.3f}')
        if loss_in_a_epoch < best_loss:
            best_loss = loss_in_a_epoch
            torch.save(model.state_dict(), '../model/RNN.pth')
    
    end_time = time.time()
    elapsed_time_minute = (end_time - start_time) / 60
    print(f"Total running time: {elapsed_time_minute:.2f} minutes")

if __name__ == '__main__':
    train()

        最后得分是0.375，是非常低的一个分数。我尝试用超参数优化，最后文件上传失败，用LSTM模型进行建模，最后文件也上传失败了，以下是我的LSTM建模代码，希望有高手指点一下哪里出问题了导致文件上传失败（或者可能是系统设置不允许用LSTM？）：

import re
import time
import pandas as pd
from typing import List, Tuple
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset
import matplotlib.pyplot as plt

# Define the LSTM model
class LSTMModel(nn.Module):
    def __init__(self, num_embed, input_size, hidden_size, output_size, num_layers, dropout, device):
        super(LSTMModel, self).__init__()
        self.embed = nn.Embedding(num_embed, input_size)
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers=num_layers, 
                            batch_first=True, dropout=dropout, bidirectional=True)
        self.fc = nn.Sequential(nn.Linear(2 * num_layers * hidden_size, output_size),
                                nn.Sigmoid(),
                                nn.Linear(output_size, 1),
                                nn.Sigmoid())

    def forward(self, x):
        # x : [bs, seq_len]
        x = self.embed(x)
        # x : [bs, seq_len, input_size]
        _, (hn, _) = self.lstm(x)  # hn : [2*num_layers, bs, h_dim]
        hn = hn.transpose(0, 1)
        z = hn.reshape(hn.shape[0], -1)  # z shape: [bs, 2*num_layers*h_dim]
        output = self.fc(z).squeeze(-1)  # output shape: [bs, 1]
        return output

# Tokenizer class
class SmilesTokenizer:
    def __init__(self, pad_token, regex, vocab_file, max_length):
        self.pad_token = pad_token
        self.regex = regex
        self.vocab_file = vocab_file
        self.max_length = max_length

        with open(self.vocab_file, "r") as f:
            lines = f.readlines()
        lines = [line.strip("\n") for line in lines]
        vocab_dic = {}
        for index, token in enumerate(lines):
            vocab_dic[token] = index
        self.vocab_dic = vocab_dic

    def _regex_match(self, smiles):
        regex_string = r"(" + self.regex + r"|"
        regex_string += r".)"
        prog = re.compile(regex_string)

        tokenized = []
        for smi in smiles:
            tokens = prog.findall(smi)
            if len(tokens) > self.max_length:
                tokens = tokens[:self.max_length]
            tokenized.append(tokens)  # Return a list of all strings
        return tokenized
    
    def tokenize(self, smiles):
        tokens = self._regex_match(smiles)
        # Add start and end tokens: <CLS>, <SEP>
        tokens = [["<CLS>"] + token + ["<SEP>"] for token in tokens]
        tokens = self._pad_seqs(tokens, self.pad_token)
        token_idx = self._pad_token_to_idx(tokens)
        return tokens, token_idx

    def _pad_seqs(self, seqs, pad_token):
        pad_length = max([len(seq) for seq in seqs])
        padded = [seq + ([pad_token] * (pad_length - len(seq))) for seq in seqs]
        return padded

    def _pad_token_to_idx(self, tokens):
        idx_list = []
        for token in tokens:
            tokens_idx = []
            for i in token:
                if i in self.vocab_dic.keys():
                    tokens_idx.append(self.vocab_dic[i])
                else:
                    self.vocab_dic[i] = max(self.vocab_dic.values()) + 1
                    tokens_idx.append(self.vocab_dic[i])
            idx_list.append(tokens_idx)
        
        return idx_list

# Read and process data
def read_data(file_path, train=True):
    df = pd.read_csv(file_path)
    reactant1 = df["Reactant1"].tolist()
    reactant2 = df["Reactant2"].tolist()
    product = df["Product"].tolist()
    additive = df["Additive"].tolist()
    solvent = df["Solvent"].tolist()
    if train:
        react_yield = df["Yield"].tolist()
    else:
        react_yield = [0 for i in range(len(reactant1))]
    
    input_data_list = []
    for react1, react2, prod, addi, sol in zip(reactant1, reactant2, product, additive, solvent):
        input_info = ".".join([react1, react2])
        input_info = ">".join([input_info, prod])
        input_data_list.append(input_info)
    output = [(react, y) for react, y in zip(input_data_list, react_yield)]

    return output

class ReactionDataset(Dataset):
    def __init__(self, data: List[Tuple[List[str], float]]):
        self.data = data
        
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

def collate_fn(batch, tokenizer):
    smi_list = [i[0] for i in batch]
    yield_list = [i[1] for i in batch]
    tokenizer_batch = torch.tensor(tokenizer.tokenize(smi_list)[1])
    yield_list = torch.tensor(yield_list)
    return tokenizer_batch, yield_list

def train():
    ## Super parameters
    N = 10  # Number of data samples
    NUM_EMBED = 294  # nn.Embedding()
    INPUT_SIZE = 300  # src length
    HIDDEN_SIZE = 512
    OUTPUT_SIZE = 512
    NUM_LAYERS = 10
    DROPOUT = 0.2
    CLIP = 1  # CLIP value
    N_EPOCHS = 100
    LR = 0.0001
    
    start_time = time.time()  # Start timer
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    data = read_data("../dataset/round1_train_data.csv")
    dataset = ReactionDataset(data)
    subset_indices = list(range(N))
    subset_dataset = Subset(dataset, subset_indices)
    tokenizer = SmilesTokenizer("<PAD>", r"\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9]", "../vocab_full.txt", max_length=300)
    train_loader = DataLoader(subset_dataset, batch_size=128, shuffle=True, collate_fn=lambda x: collate_fn(x, tokenizer))

    model = LSTMModel(NUM_EMBED, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE, NUM_LAYERS, DROPOUT, device).to(device)
    model.train()
    
    optimizer = optim.Adam(model.parameters(), lr=LR)
    criterion = nn.L1Loss()  # MAE

    best_loss = float('inf')
    for epoch in range(N_EPOCHS):
        epoch_loss = 0
        for i, (src, y) in enumerate(train_loader):
            src, y = src.to(device), y.to(device)
            optimizer.zero_grad()
            output = model(src)
            loss = criterion(output, y)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP)
            optimizer.step()
            epoch_loss += loss.item()
        
        loss_in_a_epoch = epoch_loss / len(train_loader)
        print(f'Epoch: {epoch+1:02} | Train Loss: {loss_in_a_epoch:.3f}')
        
        if loss_in_a_epoch < best_loss:
            best_loss = loss_in_a_epoch
            # Save model
            torch.save(model.state_dict(), '../model/LSTM.pth')

    end_time = time.time()  # End timer
    elapsed_time_minute = (end_time - start_time) / 60
    print(f"Total running time: {elapsed_time_minute:.2f} minutes")

if __name__ == '__main__':
    train()