<jupyter_start><jupyter_text>About This NotebookThis notebook is based on https://www.kaggle.com/konradb/model-train-efficientnet & https://www.kaggle.com/konradb/model-infer-efficientnet, with a final score of 8.90 achieved in the BMS competition.  Import Libraries<jupyter_code>import os
import re
import gc
import cv2
import timm
import time
import math
import torch
import random
import Levenshtein
import numpy as np
import pandas as pd
import torch.nn as nn
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.optim import Adam, SGD
import torchvision.models as models
from torch.nn.parameter import Parameter
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence
from sklearn.model_selection import StratifiedKFold, GroupKFold, KFold
from logging import getLogger, INFO, FileHandler,  Formatter,  StreamHandler
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, CosineAnnealingLR, ReduceLROnPlateau
from albumentations import (
    Compose, OneOf, Normalize, Resize, RandomResizedCrop, RandomCrop, HorizontalFlip, VerticalFlip, 
    RandomBrightness, RandomContrast, RandomBrightnessContrast, Rotate, ShiftScaleRotate, Cutout, 
    IAAAdditiveGaussianNoise, Transpose, Blur
    )
from albumentations.pytorch import ToTensorV2
from albumentations import ImageOnlyTransform
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)<jupyter_output><empty_output><jupyter_text>Read Input Data> Import the train dataframe containing image IDs, InChI strings, their actual lengths and parsed sequences.<jupyter_code># read the input data contained in the pickle file saved previously
train_df = pd.read_pickle('../../data/train.pkl')<jupyter_output><empty_output><jupyter_text>Add File Paths> Make the process of reading the input data more efficient by storing paths to files in the train dataframe.<jupyter_code>def get_file_path(image_id: str) -> str:
    
    """
    This method returns the paths to train images by indexing into the overall directory
    and the image_id's components.
    
    :param image_id: ID of the image
    :type  image_id: str
    :return:         path to image
    :rtype:          str
    """
    
    # index into original train images if '-' is not present
    return '../../data/bms-molecular-translation/train/{}/{}/{}/{}.png'.format(
        image_id[0], image_id[1], image_id[2], image_id
    )
# get file paths
train_df['file_path'] = train_df['image_id'].apply(get_file_path)
train_df.to_csv('../../data/train_df.csv')
# import the file back
train_df = pd.read_csv('../../data/train_df.csv')
# limit to ±300K data sub-selected by size 200-350 (HxW)
valid_ids = pd.read_csv('../../data/bmssmalldataset/new_dataset.csv')['image_id']
train_df  = train_df[train_df['image_id'].isin(valid_ids)]
print(train_df.shape)
train_df.reset_index(inplace=True)
train_len = int(len(train_df) * 0.8)

# keep 20% as test dataset
test_df   = train_df.loc[train_len:, :]

# keep 80% as train dataset
train_df  = train_df.loc[:train_len, :]

# save as csv files
test_df.to_csv('../../data/test.csv')
train_df.to_csv('../../data/reduced_train.csv')<jupyter_output><empty_output><jupyter_text>Split InChI label to Get Chemical Formula<jupyter_code>train_df['InChI_0'] = train_df['InChI'].apply(lambda x: x.split('/')[1])
train_df['InChI_0_len'] = train_df['InChI_0'].apply(lambda x: len(x))
train_df['InChI_0_text'] = train_df['InChI_text'].apply(lambda x: x.split('/')[1])<jupyter_output><empty_output><jupyter_text>Read the Tokenizer> Import the string to index mapping for each InChI token, saved previously.<jupyter_code>class Tokenizer(object):
    
    def __init__(self):
        # string to integer mapping
        self.stoi = {}
        # integer to string mapping
        self.itos = {}
    
    def __len__(self) -> None:
        
        """
        This method returns the length of token:index map.
        
        :return: length of map
        :rtype: int
        """
        # return the length of the map
        return len(self.stoi)
    
    def fit_on_texts(self, texts: list) -> None:
        
        """
        This method creates a vocabulary of all tokens contained in provided texts,
        and updates the mapping of token to index, and index to token.
        
        :param texts: list of texts
        :type texts:  list
        """
        
        # create a storage for all tokens
        vocab = set()
        
        # add tokens from each text to vocabulary
        for text in texts:
            vocab.update(text.split(' '))
            
        # sort the vocabulary in alphabetical order
        vocab = sorted(vocab)
        
        # add start, end and pad for sentence
        vocab.append('<sos>')
        vocab.append('<eos>')
        vocab.append('<pad>')
        
        # update the string to integer mapping, where integer is the index of the token
        for i, s in enumerate(vocab):
            self.stoi[s] = i
        
        # reverse the previous vocabulary to create integer to string mapping
        self.itos = {item[1]: item[0] for item in self.stoi.items()}
        
    def text_to_sequence(self, text: str) -> list:
        
        """
        This method converts the given text to a list of its individual tokens,
        including start and end of string symbols.
        
        :param text: input textual data
        :type  text: str
        :return:     list of tokens
        :rtype:      list
        """
        
        # storage to append symbols to
        sequence = []
        
        # add the start of string symbol to storage
        sequence.append(self.stoi['<sos>'])
        
        # add each token in text to storage
        for s in text.split(' '):
            sequence.append(self.stoi[s])
            
        # add the end of string symbol to storage
        sequence.append(self.stoi['<eos>'])
        
        return sequence
    
    def texts_to_sequences(self, texts: list) -> list:
        
        """
        This method converts each text in the provided list into sequences of characters.
        Each sequence is appended to a list and the said list is returned.
        
        :param texts: a list of input texts
        :type  texts: list
        :return:      a list of sequences
        :rtype:       list
        """
        
        # storage to append sequences to
        sequences = []
        
        # for each text do
        for text in texts:
            # convert the text to a list of characters
            sequence = self.text_to_sequence(text)
            # append the lists of characters to an aggregated list storage
            sequences.append(sequence)

        return sequences
    
    def sequence_to_text(self, sequence: list) -> str:
        
        """
        This method converts the sequence of characters back into text.
        
        :param sequence: list of characters
        :type  sequence: list
        :return:         text
        :rtype:          str 
        """
        # join the characters with no space in between
        return ''.join(list(map(lambda i: self.itos[i], sequence)))
    
    def sequences_to_texts(self, sequences: list) -> list:
        
        """
        This method converts each provided sequence into text and returns all texts inside a list.
        
        :param sequences: list of character sequences
        :type  sequences: list
        :return:          list of texts
        :rtype:           list
        """
        
        # storage for texts
        texts = []
        
        # convert each sequence to text and append to storage
        for sequence in sequences:
            text = self.sequence_to_text(sequence)
            texts.append(text)

        return texts
    
    def predict_caption(self, sequence: list) -> str:
        
        """
        This method predicts the caption by adding each symbol in sequence to a resulting string.
        This keeps happening up until the end of sentence or padding is met.
        
        :param sequence: list of characters
        :type  sequence: list
        :return:         image caption
        :rtype:          string
        """
        
        # storage for the final caption
        caption = ''
        
        # for each index in a sequence of symbols
        for i in sequence:
            # if symbol is the end of sentence or padding, break
            if i == self.stoi['<eos>'] or i == self.stoi['<pad>']:
                break
            # otherwise, add the symbol to the final caption
            caption += self.itos[i]
            
        return caption
    
    def predict_captions(self, sequences: list) -> list:
        
        """
        This method predicts the captions for each sequence in a list of sequences.
        
        :param sequences: list of sequences
        :type  sequences: list
        :return:          list of final image captions
        :rtype:           list
        """
        
        # storage for captions
        captions = []
        
        # for each sequence, do
        for sequence in sequences:
            
            # predict the caption per sequence
            caption = self.predict_caption(sequence)
            
            # append to the storage of captions
            captions.append(caption)
            
        return captions

# load the saved tokenizer and print its string to index mapping
tokenizer = torch.load('../../data/lstm-model/results-6/results-6/tokenizer.pth')<jupyter_output><empty_output><jupyter_text>Setup Configurations> Set configurations needed for modelling and training in a separate class.<jupyter_code>class CFG:
    
    """
    Set configurations for modelling and training.
    """
    
    debug       = False
    apex        = False
    max_len     = train_df['InChI_0_len'].max()
    
    print_freq  = 10000
    num_workers = 4
    model_name  = 'efficientnet_b1'
    enc_size    = 1280
    samp_size   = 1000
    
    size        = 288
    
    scheduler   = 'CosineAnnealingLR'
    epochs      = 50
    T_max       = 4
    
    encoder_lr  = 1e-4
    decoder_lr  = 4e-4
    min_lr      = 1e-6
    
    batch_size   = 32
    weight_decay = 1e-6
    
    gradient_accumulation_steps = 1
    max_grad_norm               = 5
    
    attention_dim = 256
    embed_dim     = 256
    decoder_dim   = 512
    dropout       = 0.5
    seed          = 42
    n_fold        = 2
    trn_fold      = [0]
    train         = True
    
    prev_model = '../../data/efficientnet_b1_fold1_best_inchi1.pth'
# if in debug mode
if CFG.debug:
    
    # set number of epochs to 1
    CFG.epochs = 1
    
    # reduce the train set to a 1000 examples
    train_df   = train_df.sample(n=CFG.samp_size, random_state=CFG.seed).reset_index(drop=True)<jupyter_output><empty_output><jupyter_text>Utilities> This is a set of utility functions used throughout the computations.<jupyter_code>def get_score(y_true: str, y_pred: str) -> float:
    
    """
    This function computes the Levenstein distance between a true label and a prediction.
    This gets computed for all the provided data and an average score is then returned.
    
    :param y_true: true InChI label
    :type  y_true: str
    :param y_pred: predicted InChI label
    :type  y_pred: str
    :return:       average Levenstein score
    :rtype:        float
    """
    
    # storage for all Levenstein scores
    scores = []
    
    # for each (true label, predicted label) pair, do
    for true, pred in zip(y_true, y_pred):
        
        # find Levenstein distance for the pair and append to storage
        score = Levenshtein.distance(true, pred)
        scores.append(score)
    
    # compute mean Levenstein distance
    avg_score = np.mean(scores)
    
    return avg_score
def init_logger(log_file: str ='./train.log'):
    
    """
    Initialize the logger file for training.
    
    :param log_file: name of the logger file
    :type  log_file: str
    :return:         logger
    :rtype:          object
    """
    
    # make a reference to a logger instance
    logger = getLogger(__name__)
    
    # specify lowest-severity log message to be handled by the logger
    logger.setLevel(INFO)
    
    # send logging outputs to stream, i.e. sys.out, and format as message
    handler1 = StreamHandler()
    handler1.setFormatter(Formatter("%(message)s"))
    
    # send logging outputs to a disk file and format as message
    handler2 = FileHandler(filename=log_file)
    handler2.setFormatter(Formatter("%(message)s"))
    
    # add both handlers to the logger
    logger.addHandler(handler1)
    logger.addHandler(handler2)
    
    return logger

# initialize the logger
LOGGER = init_logger()
def seed_torch(seed: int=42) -> None:
    
    """
    Seed torch with a specific seed number to ensure code consistency across runs.
    
    :param seed: seed number
    :type  seed: int
    """
    
    # set random seed
    random.seed(seed)
    
    # set environment seed
    os.environ['PYTHONHASHSEED'] = str(seed)
    
    # set numpy seed
    np.random.seed(seed)
    
    # set torch seeds
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    
# seed torch with the seed from configs
seed_torch(seed=CFG.seed)<jupyter_output><empty_output><jupyter_text>Cross-Validation Split> Here cross validation splits are created for the train dataframe.<jupyter_code># create a copy of the train dataframe to modify
folds = train_df.copy()

# provide train/validation indices to split data into train/validation sets
Fold  = StratifiedKFold(n_splits=CFG.n_fold, shuffle=True, random_state=CFG.seed)

# for fold_number and (train_index, validation_index) in the splitted fold
for n, (train_idx, val_idx) in enumerate(Fold.split(folds, folds['InChI_length'])): # folds = separator, 
                                                                                    # len = maxsplit
    # assign fold number
    folds.loc[val_idx, 'fold'] = int(n)

# convert fold number to integer
folds['fold'] = folds['fold'].astype(int)

# print the size of each fold
print(folds.groupby(['fold']).size())<jupyter_output><empty_output><jupyter_text>Dataset> Return the input data of the following format: (image, label tensor, label length). This is needed to ensure that the LSTM-Attention model is working as intended.<jupyter_code>class TrainDataset(Dataset):
    
    """
    This class stores train dataset attributes and methods.
    """
    
    def __init__(self, df, tokenizer, transform=None):
        
        """
        Initialize train dataset attributes.
        
        :param df:        train dataframe
        :type  df:        pd.DataFrame
        :param tokenizer: string tokenizer
        :type tokenizer:  object
        :param transform: torch transformation
        :type transform:  object
        """
        
        # inherit from parent class
        super().__init__()
        
        # assign train dataframe
        self.df = df
        
        # assign tokenizer
        self.tokenizer = tokenizer
        
        # assign file paths
        self.file_paths = df['file_path'].values
        
        # assign tokenized labels
        self.labels     = df['InChI_0_text'].values
        
        # assign transformations
        self.transform  = transform
        
    def __len__(self) -> int:
        
        """
        Return size of the train dataframe.
        
        :return: length of the dataframe
        :rtype:  int
        """
        return len(self.df)
    
    def __getitem__(self, idx: int) -> tuple:
        
        """
        Get the image, tensor of its label and label length at the inputted index.
        
        :param idx: index of dataframe
        :type  idx: int
        :return:    image, tensored label and label length
        :rtype:     tuple 
        """
        
        # get file path of the indexed item
        file_path = self.file_paths[idx]
        
        # read in the image using the file path
        image     = cv2.imread(file_path)
        
        # convert the image to RGB and float type
        image     = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        
        # if transform is specified
        if self.transform:
            
            # augment the image with the transform
            augmented = self.transform(image=image)
            
            # update the image reference to point to the transformed image
            image     = augmented['image']
        
        # get the label of the indexed item
        label = self.labels[idx]
        
        # convert label to a sequence of symbols
        label = self.tokenizer.text_to_sequence(label)
        
        # get the length of the label and convert it to Tensor
        label_length = len(label)
        label_length = torch.LongTensor([label_length])
        
        return image, torch.LongTensor(label), label_length
class TestDataset(Dataset):
    
    """
    This class stores test dataset attributes and methods.
    """
    
    def __init__(self, df, transform=None):
        
        """
        Initialize test dataset attributes.
        
        :param df:        test dataframe
        :type  df:        pd.DataFrame
        :param tokenizer: string tokenizer
        :type tokenizer:  object
        :param transform: torch transformation
        :type transform:  object
        """
        
        # inherit from parent class
        super().__init__()
        
        # assign train dataframe
        self.df = df
        
        # assign file paths
        self.file_paths = df['file_path'].values
        
        # assign transformations
        self.transform  = transform
        
        # assign fixed transformations
        self.fix_transform = Compose([Transpose(p=1), VerticalFlip(p=1)])
        
    def __len__(self) -> int:
        
        """
        Return size of the train dataframe.
        
        :return: length of the dataframe
        :rtype:  int
        """
        return len(self.df)
    
    def __getitem__(self, idx: int) -> np.array:
        
        """
        Get the image, tensor of its label and label length at the inputted index.
        
        :param idx: index of dataframe
        :type  idx: int
        :return:    image
        :rtype:     array
        """
        
        # get file path of the indexed item
        file_path = self.file_paths[idx]
        
        # read in the image using the file path
        image     = cv2.imread(file_path)
        
        # convert the image to RGB and float type
        image     = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
        
        # get image shape
        h, w, _ = image.shape
        
        # if height exceeds width, fix-transform the image
        if h > w:
            image = self.fix_transform(image=image)['image']
        
        # if transform is specified
        if self.transform:
            
            # augment the image with the transform
            augmented = self.transform(image=image)
            
            # update the image reference to point to the transformed image
            image     = augmented['image']

        
        return image
def bms_collate(batch: tuple) -> tuple:
    
    """
    Combine images, labels and label lengths per batch.
    
    :param batch: a collection of data points
    :type  batch: tuple
    :return:      stacked images, labels and label lengths, i.e. batch
    :rtype:       tuple
    """
    
    # initialize storages for images, labels and label lengths
    imgs, labels, label_lengths = [], [], []
    
    # for each data point, append image, labels and label lengths to respective storages
    for data_point in batch:
        imgs.append(data_point[0])
        labels.append(data_point[1])
        label_lengths.append(data_point[2])
    
    # pad each label sequence with the <pad> index value
    labels = pad_sequence(labels, batch_first=True, padding_value=tokenizer.stoi["<pad>"])
 
    return torch.stack(imgs), labels, torch.stack(label_lengths).reshape(-1, 1)<jupyter_output><empty_output><jupyter_text>Transformations> Define basic torch transforms for the dataset, including resizing, normalizing and tensoring.<jupyter_code>def get_transforms(*, data):
    
    """
    Compose several transforms together, mainly resizing, normalizing with Imagenet weights and tensoring.
    
    :param data: image data
    :type  data: np.array
    :return:     transformed image data
    :rtype:      np.array
    """
    if data == 'train':
        return Compose(
            [
                Resize(CFG.size, CFG.size),
                Normalize(
                    mean=[0.485, 0.456, 0.406],
                    std=[0.229, 0.224, 0.225],
                ),
                ToTensorV2(),
            ]
        )
    
    elif data == 'valid':
        return Compose([
            Resize(CFG.size, CFG.size),
            Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225],
            ),
            ToTensorV2(),
        ])
# define train dataset from the train dataframe, using the tokenizer and transformations
train_ds = TrainDataset(train_df, tokenizer, transform=get_transforms(data='train'))
# display an example from the train dataset
for i in range(1):
    
    # get image, label and label length from the train dataset
    image, label, label_length = train_ds[i]
    # convert label sequence to text
    text = tokenizer.sequence_to_text(label.numpy())
    
    # transpose the image and show it
    plt.imshow(image.transpose(0,1).transpose(1,2))
    plt.title(f'label: {label} text: {text} label_length: {label_length}')
    plt.show()<jupyter_output><empty_output><jupyter_text>Modeling> Create encoder, attention and decoder classes for modeling.<jupyter_code>class Encoder(nn.Module):
    
    """
    Encodes the image with 3 color channels into a smaller learned image.
    """
    
    def __init__(self, model_name=CFG.model_name, pretrained=False):
        
        """
        Initialize the encoder with CNN equal to the chosen model and set pretrained parameter.
        
        :param model_name: name of the model
        :type  model_name: str
        :param pretrained: pretrained weights or not
        :type  pretrained: Boolean
        """
        
        # inherit attributes and methods from parent class
        super().__init__()
        
        self.cnn = timm.create_model(model_name, pretrained=pretrained)
        self.n_features = self.cnn.classifier.in_features
        
        # replacing the last two layers with no-op
        self.cnn.global_pool = nn.Identity()
        self.cnn.classifier = nn.Identity()
        
    def forward(self, x):
        
        """
        Propagate the input forward.
        
        :param x: image data
        :type  x: np.array
        :return:  image features
        :rtype:   np.array 
        """
        
        # get image features from the CNN
        features = self.cnn(x)
        
        # re-arrange the dimensions so that (bs, encoded_image_size, encoded_image_size, n_channels=2048)
        features = features.permute(0, 2, 3, 1)
        
        return features
class Attention(nn.Module):
    
    """
    Define Attention network to calculate the attention value.
    """
    
    def __init__(self, encoder_dim, decoder_dim, attention_dim):
        
        """
        Initialize the Attention network.
        
        :param encoder_dim:   input size of the encoder network
        :type  encoder_dim:   int
        :param decoder_dim:   input size of the decoder network
        :type  decoder_dim:   int
        :param attention_dim: input size of the attention network
        :type  attention_dim: int
        """
        
        # initiliaze and inherit from parent class
        super(Attention, self).__init__()
        
        # 1st linear layer to transform the encoded image
        self.encoder_att = nn.Linear(encoder_dim, attention_dim)
        
        # 1st linear layer to transform the decoder's output
        self.decoder_att = nn.Linear(decoder_dim, attention_dim)
        
        # linear layer to calculate values to be softmax-ed
        self.full_att    = nn.Linear(attention_dim, 1)
        
        # define ReLU function
        self.relu        = nn.ReLU()
        
        # define the softmax layer to calculate weights
        self.softmax     = nn.Softmax(dim=1)
        
    def forward(self, encoder_out, decoder_hidden):
        
        """
        Propagate the inputs forward.
        
        :param encoder_out: encoded images of dimension (batch_size, num_pixels, encoder_dim)
        :type  encoder_out: tensor
        :param decoder_hidden: previous decoder output of dimension (batch_size, decoder_dim)
        :type  decoder_hidden: tensor
        :return: attention-weighted encoding, weights
        :rtype:  tensor of (batch_size, encoder_dim)
        """
        
        # apply the linear layer to encoded images to get (bs, num_pixels, attention_dim)
        att1 = self.encoder_att(encoder_out)
        
        # apply the linear layer to transform the decoder's output to get (bs, attention_dim)
        att2 = self.decoder_att(decoder_hidden)
        
        # pass the sum of transformed encoded images and decoder outputs through a ReLU and apply a linear layer
        # gets (bs, num_pixels)
        att  = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2)
        
        # pass the results through a softmax layer to get weights
        alpha = self.softmax(att)
        
        # apply the resulting (bs, num_pixels) weights to the encoder output and sum across dim=1
        attention_weighted_encoding = (encoder_out * alpha.unsqueeze(2)).sum(dim=1)
        
        return attention_weighted_encoding, alpha
class DecoderWithAttention(nn.Module):
    
    """
    Decoder network with attention network used for training.
    """
    
    def __init__(self, attention_dim, embed_dim, decoder_dim, vocab_size, device, encoder_dim=CFG.enc_size, dropout=CFG.dropout):
        
        """
        Initialize the decoder with attention network.
        
        :param attention_dim: input size of the attention network
        :type  attention_dim: int
        :param embed_dim:     input size of the embedding network
        :type  embed_dim:     int
        :param decoder_dim:   input size of the decoder network
        :type  decoder_dim:   int
        :param vocab_size:    total number of characters used in training
        :type  vocab_size:    int
        :param encoder_dim:   input size of the encoder network
        :type  encoder_dim:   int
        :param dropout:       dropout rate
        :type  dropout:       float
        """
        
        # inherit from parent class
        super(DecoderWithAttention, self).__init__()
        
        # set dimensions of the encoder, attention, embedder and decoder
        self.encoder_dim   = encoder_dim
        self.attention_dim = attention_dim
        self.embed_dim     = embed_dim
        self.decoder_dim   = decoder_dim
        
        # set vocabulary size, dropout rate and the used device
        self.vocab_size    = vocab_size
        self.dropout       = dropout
        self.device        = device
        
        # set attention network, embedding network and dropout layer
        self.attention     = Attention(encoder_dim, decoder_dim, attention_dim)
        self.embedding     = nn.Embedding(vocab_size, embed_dim)
        self.dropout       = nn.Dropout(p=self.dropout)
        
        # set the LSTM Cell decoder
        self.decode_step   = nn.LSTMCell(embed_dim+encoder_dim, decoder_dim, bias=True)
        
        # linear layer to find the initial hidden state of LSTM Cell
        self.init_h        = nn.Linear(encoder_dim, decoder_dim)
        
        # linear layer to find the initial cell state of LSTM Cell
        self.init_c        = nn.Linear(encoder_dim, decoder_dim)
        
        # linear layer to create a sigmoid-activated gate
        self.f_beta        = nn.Linear(decoder_dim, encoder_dim)
        self.sigmoid       = nn.Sigmoid()
        
        # linear layer to find scores over vocabulary
        self.fc            = nn.Linear(decoder_dim, vocab_size)
        
        # initialize some layers with uniform distribution
        self.init_weights()
        
    def init_weights(self):
        
        """
        Initialize weights with uniform distribution for embedding and FC layers.
        """
        
        self.embedding.weight.data.uniform_(-0.1, 0.1)
        self.fc.bias.data.fill_(0)
        self.fc.weight.data.uniform_(-0.1, 0.1)
        
    def load_pretrained_embeddings(self, embeddings):
        
        """
        Load the embedding layer with pre-trained embeddings.
        
        :param embeddings: pre-trained embeddings
        """
        
        self.embedding.weight = nn.Parameter(embeddings)
        
    def fine_tune_embeddings(self, fine_tune=True):
        
        """
        Allow fine-tuning of the embedding layer.
        
        :param fine_tune: allow fine-tuning
        :type  fine_tune: Boolean
        """
        # loop over each embedding parameter to set the fine-tuning option
        for p in self.embedding.parameters():
            p.requires_grad = fine_tune
            
    def init_hidden_state(self, encoder_out):
        
        """
        Create the initial hidden and cell state for decoder's LSTM based on encoded images.
        
        :param encoder_out: encoded images, of size (bs, num_pixels, encoder_dim)
        :type  encoder_out: tensor
        :return:            initial hidden and cell states
        :rtype:             tuple
        
        """
        
        # get the mean of the encoded image's dim=1
        mean_encoder_out = encoder_out.mean(dim=1)
        
        # initialize the hidden and cell states
        h                = self.init_h(mean_encoder_out)
        c                = self.init_c(mean_encoder_out)
        
        return h, c
    
    def forward(self, encoder_out, encoded_captions, caption_lengths):
        
        """
        Propagate the inputs forward.
        
        :param encoder_out: encoded images of dimension (bs, num_pixels, encoder_dim)
        :type  encoder_out: tensor
        :param encoded_captions: encoded captions of dimension (bs, max_caption_length)
        :type  encoded_captions: tensor
        :param caption_lengths: caption lengths of dimension (bs, 1)
        :type  caption_lengths: tensor
        :return: scores for vocabulary, sorted encoded captions, decode lengths, weights, sort indices
        :rtype:  tuple
        """
        batch_size  = encoder_out.size(0)
        encoder_dim = encoder_out.size(-1)
        vocab_size  = self.vocab_size
        
        # flatten the image
        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)
        num_pixels  = encoder_out.size(1)
        
        # sort input data by decreasing lengths
        caption_lengths, sort_ind = caption_lengths.squeeze(1).sort(dim=0, descending=True)
        encoder_out = encoder_out[sort_ind]
        encoded_captions = encoded_captions[sort_ind]
        
        # embed the encoded captions
        embeddings = self.embedding(encoded_captions)
        
        # initialize LSTM state
        h, c = self.init_hidden_state(encoder_out)
        
        # decoding lengths are actual lengths - 1 because of <end>
        decode_lengths = (caption_lengths - 1).tolist()
        
        # initialize tensors to hold word predictions scores and alphas (weights)
        predictions    = torch.zeros(batch_size, max(decode_lengths), vocab_size).to(self.device)
        alphas         = torch.zeros(batch_size, max(decode_lengths), num_pixels).to(self.device)
        
        # at each time step, decode by 
        for t in range(max(decode_lengths)):
            batch_size_t = sum([l>t for l in decode_lengths])
            
            # attention weighted encoder's output based on decoder's previous state output
            attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t], h[:batch_size_t])
            gate = self.sigmoid(self.f_beta(h[:batch_size_t]))
            attention_weighted_encoding = gate * attention_weighted_encoding
            
            # generate a new word with previous word and attention weighted encoding
            h, c = self.decode_step(
                torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),
                (h[:batch_size_t], c[:batch_size_t])
            )
            preds = self.fc(self.dropout(h))
            predictions[:batch_size_t, t, :] = preds
            alphas[:batch_size_t, t, :] = alpha
            
        return predictions, encoded_captions, decode_lengths, alphas, sort_ind
    
    def predict(self, encoder_out, decode_lengths, tokenizer):
        
        """
        Predict captions.
        
        :param encoder_out: encoded images of dimension (bs, num_pixels, encoder_dim)
        :type  encoder_out: tensor
        :param decode_lengths: lengths of decoded captions
        :type  decode_lengths: tensor
        :param tokenizer: word tokenizer
        :type  tokenizer: dict
        :return: predictions
        :rtype:  str
        """
        batch_size = encoder_out.size(0)
        encoder_dim = encoder_out.size(-1)
        vocab_size = self.vocab_size
        
        # flatten the image
        encoder_out = encoder_out.view(batch_size, -1, encoder_dim)
        num_pixels = encoder_out.size(1)
        
        # define start tokens and embed them
        start_tokens = torch.ones(batch_size, dtype=torch.long).to(self.device) * tokenizer.stoi["<sos>"]
        embeddings = self.embedding(start_tokens)
        
        # initialize LSTM state
        h, c = self.init_hidden_state(encoder_out)  # (batch_size, decoder_dim)
        
        # initialize tensors to hold predictions and final conditions
        predictions = torch.zeros(batch_size, decode_lengths, vocab_size).to(self.device)
        end_condition = torch.zeros(batch_size, dtype=torch.long).to(encoder_out.device)
        
        # at each time step decode by
        for t in range(decode_lengths):
            
            # applying attention to encoded image and decoder's previous state output
            attention_weighted_encoding, alpha = self.attention(encoder_out, h)
            gate = self.sigmoid(self.f_beta(h))  # gating scalar, (batch_size_t, encoder_dim)
            attention_weighted_encoding = gate * attention_weighted_encoding
            
            # generate a new word with previous word and attention weighted encoding
            h, c = self.decode_step(
                torch.cat([embeddings, attention_weighted_encoding], dim=1),
                (h, c))  # (batch_size_t, decoder_dim)
            preds = self.fc(self.dropout(h))  # (batch_size_t, vocab_size)
            predictions[:, t, :] = preds
            
            # if end of sentence is reached, stop predicting
            end_condition |= (torch.argmax(preds, -1) == tokenizer.stoi["<eos>"])
            if end_condition.sum() == batch_size:
                break
                
            # embed the predictions
            embeddings = self.embedding(torch.argmax(preds, -1))
            
        return predictions<jupyter_output><empty_output><jupyter_text>Helper Functions> This is a set of functions used as utilities.<jupyter_code>class AverageMeter(object):
    
    """
    Computes and stores average and current values.
    """
    
    def __init__(self):
        
        """
        Reset settings.
        """
        self.reset()
    
    def reset(self):
        
        """
        Set current, average values, sum and count values to zero.
        """
        
        self.val   = 0
        self.avg   = 0
        self.sum   = 0
        self.count = 0
        
    def update(self, val, n=1):
        
        """
        Update current value, sum, count and average value.
        """
        
        self.val    = val
        self.sum   += val*n
        self.count += n
        self.avg = self.sum / self.count
def asMinutes(s):
    
    """
    Convert seconds to minutes.
    
    :param s: seconds
    :type  s: float
    """
    m  = math.floor(s/60)
    s -= m * 60
    return '%d %ds'% (m, s)
def timeSince(since, percent):
    
    """
    Calculate time since.
    
    :param since: previous date
    :type  since: time
    """
    
    now = time.time()
    s   = now - since
    es  = s / (percent)
    rs  = es - s
    
    return '%s (remain %s)' % (asMinutes(s), asMinutes(rs))
def train_fn(
    train_loader, encoder, decoder, 
    criterion, encoder_optimizer, decoder_optimizer, 
    epoch, encoder_scheduler, decoder_scheduler, device
):
    """
    Perform one epoch training.
    
    :param train_loader: data loader for training data
    :type  train_loader: DataLoader
    :param encoder:      encoder model
    :type  encoder:      Encoder
    :param decoder:      decoder model
    :type  decoder:      Decoder
    :param criterion:    loss layer
    :type criterion:     Loss 
    :param encoder_optimizer: optimizer for encoder
    :type  encoder_optimizer: Optimizer
    :param decoder_optimizer: optimizer for decoder
    :type  decoder_optimizer: Optimizer
    :param epoch:             Epoch number
    :type  epoch:             int
    :param encoder_scheduler: Encoder scheduler
    :type  encoder_scheduler: Encoder
    :param decoder_scheduler: Decoder scheduler
    :type  decoder_scheduler: Decoder
    :param device:            device selection
    :type  device:            Device
    :return:                  Average loss
    :rtype:                   float
    """
    
    batch_time = AverageMeter()
    data_time  = AverageMeter()
    losses     = AverageMeter()
    
    # switch to train mode
    encoder.train()
    decoder.train()
    start = end = time.time()
    global_step = 0
    
    # for index and inputs (imgs, labels and label lengths) in train dataset, do
    for step, (images, labels, label_lengths) in enumerate(train_loader):
        
        # update time step
        data_time.update(time.time() - end)
        
        # send inputs to device
        images = images.to(device)
        labels = labels.to(device)
        label_lengths = label_lengths.to(device)
        
        # set batch size
        batch_size = images.size(0)
        
        # encode images to get features
        features = encoder(images)
        
        # return scores for vocabulary, sorted encoded captions, decode lengths, weights, sort indices
        predictions, caps_sorted, decode_lengths, alphas, sort_ind = decoder(features, labels, label_lengths)
        
        # specify true targets (true y's) after <start>
        targets = caps_sorted[:, 1:]
        
        # get predictions in a packed sequence
        predictions = pack_padded_sequence(predictions, decode_lengths, batch_first=True).data
        
        # get targets in a packed sequence
        targets = pack_padded_sequence(targets, decode_lengths, batch_first=True).data
        
        # compute loss between predictions and targets
        loss = criterion(predictions, targets)
        
        # record the loss
        losses.update(loss.item(), batch_size)
        
        # if optimizer steps are more than 1, divide loss by the number of those steps
        if CFG.gradient_accumulation_steps > 1:
            loss = loss / CFG.gradient_accumulation_steps
            
        if CFG.apex:
            with amp.scale_loss(loss, decoder_optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            # backward propagation of loss
            loss.backward()
        
        # perform gradient clipping to avoid exploding gradients
        encoder_grad_norm = torch.nn.utils.clip_grad_norm_(encoder.parameters(), CFG.max_grad_norm)
        decoder_grad_norm = torch.nn.utils.clip_grad_norm_(decoder.parameters(), CFG.max_grad_norm)
        
        # update weights
        if (step + 1) % CFG.gradient_accumulation_steps == 0:
            
            # take a step based on gradients
            encoder_optimizer.step()
            decoder_optimizer.step()
            
            # clear all gradients from last step
            encoder_optimizer.zero_grad()
            decoder_optimizer.zero_grad()
            global_step += 1
            
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        
        # print out the results per epoch
        if step % CFG.print_freq == 0 or step == (len(train_loader)-1):
            print('Epoch: [{0}][{1}/{2}] '
                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
                  'Elapsed {remain:s} '
                  'Loss: {loss.val:.4f}({loss.avg:.4f}) '
                  'Encoder Grad: {encoder_grad_norm:.4f}  '
                  'Decoder Grad: {decoder_grad_norm:.4f}  '
                  .format(
                   epoch+1, step, len(train_loader), batch_time=batch_time,
                   data_time=data_time, loss=losses,
                   remain=timeSince(start, float(step+1)/len(train_loader)),
                   encoder_grad_norm=encoder_grad_norm,
                   decoder_grad_norm=decoder_grad_norm,
                   ))
            
    return losses.avg


def valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device):
    
    """
    Predict for validation set.
    
    :param valid_loader: Data Loader for validation data
    :type  valid_loader: DataLoader
    :param encoder:      Encoder
    :type  encoder:      Encoder
    :param decoder:      Decoder
    :type  decoder:      Decoder
    :param tokenizer:    Tokenizer
    :type  tokenizer:    Tokenizer
    :param criterion:    Loss
    :type  criterion:    Loss 
    :param device:       device selection
    :type  device:       Device
    :return:             predictions
    :rtype:              list
    """
    
    batch_time = AverageMeter()
    data_time = AverageMeter()
    
    # switch to evaluation mode
    encoder.eval()
    decoder.eval()
    
    # store predictions here
    text_preds = []
    
    start = end = time.time()
    
    # for each image in validation set
    for step, (images) in enumerate(valid_loader):
        
        # measure data loading time
        data_time.update(time.time() - end)
        
        # send images to device
        images = images.to(device)
        
        # specify batch size
        batch_size = images.size(0)
        
        # disable gradient calculation to avoid CUDA errors
        with torch.no_grad():
            
            # encode images
            features = encoder(images)
            
            # predict sequence using decoder
            predictions = decoder.predict(features, CFG.max_len, tokenizer)
            
        # choose the best predicted sequence
        predicted_sequence = torch.argmax(predictions.detach().cpu(), -1).numpy()
        
        # predict captions from the predicted sequence and append to storage
        _text_preds = tokenizer.predict_captions(predicted_sequence)
        text_preds.append(_text_preds)
        
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()
        
        # print out the validation results
        if step % CFG.print_freq == 0 or step == (len(valid_loader)-1):
            print('EVAL: [{0}/{1}] '
                  'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
                  'Elapsed {remain:s} '
                  .format(
                   step, len(valid_loader), batch_time=batch_time,
                   data_time=data_time,
                   remain=timeSince(start, float(step+1)/len(valid_loader)),
                   ))
            
    # concatenate preds into one string
    text_preds = np.concatenate(text_preds)
    
    return text_preds<jupyter_output><empty_output><jupyter_text>Train Loop> Performs training.<jupyter_code>def train_loop(folds, fold):
    
    """
    Perform training in a loop.
    """

    LOGGER.info(f"========== fold: {fold} training ==========")

    # ====================================================
    # loader
    # ====================================================
    trn_idx = folds[folds['fold'] != fold].index
    val_idx = folds[folds['fold'] == fold].index

    train_folds = folds.loc[trn_idx].reset_index(drop=True)
    valid_folds = folds.loc[val_idx].reset_index(drop=True)
    valid_labels = valid_folds['InChI'].values

    train_dataset = TrainDataset(train_folds, tokenizer, transform=get_transforms(data='train'))
    valid_dataset = TestDataset(valid_folds, transform=get_transforms(data='valid'))

    train_loader = DataLoader(train_dataset, 
                              batch_size=CFG.batch_size, 
                              shuffle=True, 
                              num_workers=CFG.num_workers, 
                              pin_memory=True,
                              drop_last=True, 
                              collate_fn=bms_collate)
    
    valid_loader = DataLoader(valid_dataset, 
                              batch_size=CFG.batch_size, 
                              shuffle=False, 
                              num_workers=CFG.num_workers,
                              pin_memory=True, 
                              drop_last=False)
    
    # ====================================================
    # scheduler 
    # ====================================================
    def get_scheduler(optimizer):
        if CFG.scheduler=='ReduceLROnPlateau':
            scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=CFG.factor, patience=CFG.patience, verbose=True, eps=CFG.eps)
        elif CFG.scheduler=='CosineAnnealingLR':
            scheduler = CosineAnnealingLR(optimizer, T_max=CFG.T_max, eta_min=CFG.min_lr, last_epoch=-1)
        elif CFG.scheduler=='CosineAnnealingWarmRestarts':
            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=CFG.T_0, T_mult=1, eta_min=CFG.min_lr, last_epoch=-1)
        return scheduler

    # ====================================================
    # model & optimizer
    # ====================================================
#     states = torch.load(CFG.prev_model, map_location=torch.device('cpu'))
    encoder = Encoder(CFG.model_name, pretrained=True)
#     encoder.load_state_dict(states['encoder'])
    encoder.to(device)
    encoder_optimizer = Adam(encoder.parameters(), lr=CFG.encoder_lr, weight_decay=CFG.weight_decay, amsgrad=False)
#     encoder_optimizer.load_state_dict(states['encoder_optimizer'])
    encoder_scheduler = get_scheduler(encoder_optimizer)
#     encoder_scheduler.load_state_dict(states['encoder_scheduler'])
    
    decoder = DecoderWithAttention(attention_dim=CFG.attention_dim,
                                   embed_dim=CFG.embed_dim,
                                   decoder_dim=CFG.decoder_dim,
                                   vocab_size=len(tokenizer),
                                   dropout=CFG.dropout,
                                   device=device)
    
#     decoder.load_state_dict(states['decoder'])
    decoder.to(device)
    decoder_optimizer = Adam(decoder.parameters(), lr=CFG.decoder_lr, weight_decay=CFG.weight_decay, amsgrad=False)
#     decoder_optimizer.load_state_dict(states['decoder_optimizer'])
    
    decoder_scheduler = get_scheduler(decoder_optimizer)
#     decoder_scheduler.load_state_dict(states['decoder_scheduler'])
    
    # ====================================================
    # loop
    # ====================================================
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.stoi["<pad>"])

    best_score = np.inf
    best_loss = np.inf
    
    for epoch in range(CFG.epochs):
        
        start_time = time.time()
        
        # train
        avg_loss = train_fn(train_loader, encoder, decoder, criterion, 
                            encoder_optimizer, decoder_optimizer, epoch, 
                            encoder_scheduler, decoder_scheduler, device)

        # eval
        text_preds = valid_fn(valid_loader, encoder, decoder, tokenizer, criterion, device)

        LOGGER.info(f"labels: {valid_labels[:5]}")
        LOGGER.info(f"preds: {text_preds[:5]}")
        
        # scoring
        score = get_score(valid_labels, text_preds)
        
        if isinstance(encoder_scheduler, ReduceLROnPlateau):
            encoder_scheduler.step(score)
        elif isinstance(encoder_scheduler, CosineAnnealingLR):
            encoder_scheduler.step()
        elif isinstance(encoder_scheduler, CosineAnnealingWarmRestarts):
            encoder_scheduler.step()
            
        if isinstance(decoder_scheduler, ReduceLROnPlateau):
            decoder_scheduler.step(score)
        elif isinstance(decoder_scheduler, CosineAnnealingLR):
            decoder_scheduler.step()
        elif isinstance(decoder_scheduler, CosineAnnealingWarmRestarts):
            decoder_scheduler.step()

        elapsed = time.time() - start_time

        LOGGER.info(f'Epoch {epoch+1} - avg_train_loss: {avg_loss:.4f}  time: {elapsed:.0f}s')
        LOGGER.info(f'Epoch {epoch+1} - Score: {score:.4f}')
        
        if score < best_score:
            best_score = score
            LOGGER.info(f'Epoch {epoch+1} - Save Best Score: {best_score:.4f} Model')
            torch.save({'encoder': encoder.state_dict(), 
                        'encoder_optimizer': encoder_optimizer.state_dict(), 
                        'encoder_scheduler': encoder_scheduler.state_dict(), 
                        'decoder': decoder.state_dict(), 
                        'decoder_optimizer': decoder_optimizer.state_dict(), 
                        'decoder_scheduler': decoder_scheduler.state_dict(), 
                        'text_preds': text_preds,
                       },
                        '../../data/'+f'{CFG.model_name}_fold{fold}_best_inchi1.pth')
            
        np.save('../../data/inchi1-predictions.npy', text_preds)
        np.save('../../data/inchi1-validations.npy', valid_labels)<jupyter_output><empty_output><jupyter_text>Main Program - Train<jupyter_code>def main():

    """
    Prepare: 1.train  2.folds
    """

    if CFG.train:
        # train
        for fold in range(CFG.n_fold):
            if fold in CFG.trn_fold:
                print('fold: ', fold)
                train_loop(folds, fold)
if __name__ == '__main__':
    main()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09131433733283585, 'fraction_numerical': 0.005537772715668755, 'mean_word_length': 2.057029655420819, 'pattern_counts': {'":': 0, '<': 72, '<?xml version=': 0, '>': 101, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '14', 'dir': 'jupyter-structured-clean-dedup', 'id': '67312', 'license': 'MIT', 'n_tokens_mistral': 14851, 'n_tokens_neox': 14247, 'n_words': 3925, 'path': 'source/notebooks/splitted-labels/image-to-sequence-baseline-InChI-1.ipynb', 'repo_name': 'purrwhite/fuzzy-blue', 'size': '50923'}