ReLSTM copy.py

from typing import List
import torch
import torch.nn as nn
import torch.optim as optim
import give_valid_test
import os
import math

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
n_step = 5 # number of cells(= number of Step)
n_hidden = 128 # number of hidden units in one cell
batch_size = 128 # batch size
learn_rate = 0.0005
all_epoch = 2 #the all epoch for training
emb_size = 256 #embeding size
save_checkpoint_epoch = 5 # save a checkpoint per save_checkpoint_epoch epochs !!! Note the save path !!!
data_root = 'penn_small'
train_path = os.path.join(data_root, 'train.txt') # the path of train dataset
tensors:List[List[torch.Tensor]] = list();
cells:List[List[torch.Tensor]] = list();

def make_batch(train_path, word2number_dict, batch_size, n_step):
    all_input_batch = []
    all_target_batch = []

    text = open(train_path, 'r', encoding='utf-8') #open the file

    input_batch = []
    target_batch = []
    for sen in text:
        word = sen.strip().split(" ")  # space tokenizer
        word = ["<sos>"] + word
        word = word + ["<eos>"]

        if len(word) <= n_step:   #pad the sentence
            word = ["<pad>"]*(n_step+1-len(word)) + word

        for word_index in range(len(word)-n_step):
            input = [word2number_dict[n] for n in word[word_index:word_index+n_step]]  # create (1~n-1) as input
            target = word2number_dict[word[word_index+n_step]]  # create (n) as target, We usually call this 'casual language model'
            input_batch.append(input)
            target_batch.append(target)

            if len(input_batch) == batch_size:
                all_input_batch.append(input_batch)
                all_target_batch.append(target_batch)
                input_batch = []
                target_batch = []

    return all_input_batch, all_target_batch # (batch num, batch size, n_step) (batch num, batch size)

def make_dict(train_path):
    text = open(train_path, 'r', encoding='utf-8')  #open the train file
    word_list = set()  # a set for making dict

    for line in text:
        line = line.strip().split(" ")
        word_list = word_list.union(set(line))

    word_list = list(sorted(word_list))   #set to list

    word2number_dict = {w: i+2 for i, w in enumerate(word_list)}
    number2word_dict = {i+2: w for i, w in enumerate(word_list)}

    #add the <pad> and <unk_word>
    word2number_dict["<pad>"] = 0
    number2word_dict[0] = "<pad>"
    word2number_dict["<unk_word>"] = 1
    number2word_dict[1] = "<unk_word>"
    word2number_dict["<sos>"] = 2
    number2word_dict[2] = "<sos>"
    word2number_dict["<eos>"] = 3
    number2word_dict[3] = "<eos>"

    return word2number_dict, number2word_dict


class TextLSTM(nn.Module):
    def __init__(self,n_class,emb_size,n_hidden):
        super(TextLSTM,self).__init__()
        self.C=nn.Embedding(n_class,embedding_dim=emb_size);
        self.W=nn.Linear(n_hidden,n_class,bias=False)
        self.b=nn.Parameter(torch.ones([n_class]))

        self.W_ii=nn.Linear(emb_size,n_hidden,bias=False)
        self.W_if=nn.Linear(emb_size,n_hidden,bias=False)
        self.W_ig=nn.Linear(emb_size,n_hidden,bias=False)
        self.W_io=nn.Linear(emb_size,n_hidden,bias=False)
        self.b_ii=nn.Parameter(torch.ones([n_hidden]))
        self.b_if=nn.Parameter(torch.ones([n_hidden]))
        self.b_ig=nn.Parameter(torch.ones([n_hidden]))
        self.b_io=nn.Parameter(torch.ones([n_hidden]))

        self.W_hi=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hf=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hg=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_ho=nn.Linear(n_hidden,n_hidden,bias=False)
        self.b_hi=nn.Parameter(torch.ones([n_hidden]))
        self.b_hf=nn.Parameter(torch.ones([n_hidden]))
        self.b_hg=nn.Parameter(torch.ones([n_hidden]))
        self.b_ho=nn.Parameter(torch.ones([n_hidden]))
        self.n_hidden=n_hidden
    def forward(self,X):
        X=self.C(X);
        hidden_state = torch.zeros(1, len(X), self.n_hidden)  # [num_layers(=1) * num_directions(=1), batch_size, n_hidden]
        cell_state = torch.zeros(1, len(X), self.n_hidden)
        X=X.transpose(0,1);

        for X_t in X:
            input_gate = torch.sigmoid(self.W_ii(X_t)+self.b_ii+self.W_hi(hidden_state)+self.b_hi)
            forget_gate = torch.sigmoid(self.W_if(X_t)+self.b_if+self.W_hf(hidden_state)+self.b_hf)
            cell_gate = torch.tanh(self.W_ig(X_t)+self.b_ig+self.W_hg(hidden_state)+self.b_hg)
            output_gate = torch.sigmoid(self.W_io(X_t)+self.b_io+self.W_ho(hidden_state)+self.b_ho)
            cell_state = forget_gate*cell_state+input_gate*cell_gate;
            hidden_state = output_gate*torch.tanh(cell_state);
        outputs=output_gate[-1]
        model=self.W(outputs)+self.b;
        return model;
    
class TextLSTM_1(nn.Module):
    def __init__(self,n_class,emb_size,n_hidden):
        super(TextLSTM_1,self).__init__()
        self.C=nn.Embedding(n_class,embedding_dim=emb_size);
        self.W=nn.Linear(n_hidden,n_class,bias=False)
        self.b=nn.Parameter(torch.ones([n_class]))

        self.W_ii=nn.Linear(emb_size,n_hidden,bias=False)
        self.W_if=nn.Linear(emb_size,n_hidden,bias=False)
        self.W_ig=nn.Linear(emb_size,n_hidden,bias=False)
        self.W_io=nn.Linear(emb_size,n_hidden,bias=False)
        self.b_ii=nn.Parameter(torch.ones([n_hidden]))
        self.b_if=nn.Parameter(torch.ones([n_hidden]))
        self.b_ig=nn.Parameter(torch.ones([n_hidden]))
        self.b_io=nn.Parameter(torch.ones([n_hidden]))

        self.W_hi=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hf=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hg=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_ho=nn.Linear(n_hidden,n_hidden,bias=False)
        self.b_hi=nn.Parameter(torch.ones([n_hidden]))
        self.b_hf=nn.Parameter(torch.ones([n_hidden]))
        self.b_hg=nn.Parameter(torch.ones([n_hidden]))
        self.b_ho=nn.Parameter(torch.ones([n_hidden]))
        self.n_hidden=n_hidden
    def forward(self,X):
        X=self.C(X);
        hidden_s= [] if len(tensors)==0 else tensors[0]
        cell_s=[] if len(cells)==0 else cells[0]
        hidden_state = torch.zeros(1, len(X), self.n_hidden) if len(hidden_s)==0 else hidden_s[0]  # [num_layers(=1) * num_directions(=1), batch_size, n_hidden]
        cell_state = torch.zeros(1, len(X), self.n_hidden) if len(cell_s)==0 else cell_s[0]
        length=len(X)
        X=X.transpose(0,1);
        i=0
        for X_t in X:
            if i>0: 
                hidden_state = torch.zeros(1,length,self.n_hidden) if len(hidden_s)<=i else hidden_s[i]
                cell_state = torch.zeros(1, length, self.n_hidden) if len(cell_s)<=i else cell_s[i]
            input_gate = torch.sigmoid(self.W_ii(X_t)+self.b_ii+self.W_hi(hidden_state)+self.b_hi)
            forget_gate = torch.sigmoid(self.W_if(X_t)+self.b_if+self.W_hf(hidden_state)+self.b_hf)
            cell_gate = torch.tanh(self.W_ig(X_t)+self.b_ig+self.W_hg(hidden_state)+self.b_hg)
            output_gate = torch.sigmoid(self.W_io(X_t)+self.b_io+self.W_ho(hidden_state)+self.b_ho)
            cell_state = forget_gate*cell_state+input_gate*cell_gate;
            hidden_state = output_gate*torch.tanh(cell_state);
            if len(cell_s)<=i:
                hidden_s.append(hidden_state.detach())
                cell_s.append(cell_state.detach())
            else:
                hidden_s[i]=hidden_state.detach()
                cell_s[i]=cell_state.detach()
            i=i+1
        outputs=output_gate[-1]
        model=self.W(outputs)+self.b;
        if len(tensors)==0: 
            tensors.append(hidden_s);
            cells.append(cell_s)
        else: 
            tensors[0]=hidden_s
            cells[0]=cell_s
        return torch.cat([hidden_state[0],model],dim=1);
class TextLSTM_2(nn.Module):
    def __init__(self,n_class,emb_size,n_hidden,numero:int):
        super(TextLSTM_2,self).__init__()
        self.numero=numero
        # self.C=nn.Embedding(n_class+n_hidden,embedding_dim=emb_size);
        self.W=nn.Linear(n_hidden,n_class,bias=False)
        self.b=nn.Parameter(torch.ones([n_class]))

        self.W_ii=nn.Linear(n_class+n_hidden,n_hidden,bias=False)
        self.W_if=nn.Linear(n_class+n_hidden,n_hidden,bias=False)
        self.W_ig=nn.Linear(n_class+n_hidden,n_hidden,bias=False)
        self.W_io=nn.Linear(n_class+n_hidden,n_hidden,bias=False)
        self.b_ii=nn.Parameter(torch.ones([n_hidden]))
        self.b_if=nn.Parameter(torch.ones([n_hidden]))
        self.b_ig=nn.Parameter(torch.ones([n_hidden]))
        self.b_io=nn.Parameter(torch.ones([n_hidden]))

        self.W_hi=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hf=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hg=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_ho=nn.Linear(n_hidden,n_hidden,bias=False)
        self.b_hi=nn.Parameter(torch.ones([n_hidden]))
        self.b_hf=nn.Parameter(torch.ones([n_hidden]))
        self.b_hg=nn.Parameter(torch.ones([n_hidden]))
        self.b_ho=nn.Parameter(torch.ones([n_hidden]))
        self.n_class=n_class
        self.n_hidden=n_hidden    
    def forward(self,X):
        # X=self.C(X);
        #print(X.shape)
        hidden_s= [] if len(tensors)<=self.numero+1 else tensors[self.numero+1]
        cell_s=[] if len(cells)<=self.numero+1 else cells[self.numero+1]
        hidden_state = torch.zeros(1, len(X), self.n_hidden) if len(hidden_s)==0 else hidden_s[0]  # [num_layers(=1) * num_directions(=1), batch_size, n_hidden]
        cell_state = torch.zeros(1, len(X), self.n_hidden) if len(cell_s)==0 else cell_s[0]
        
        #X=X.transpose(0,1);
        #print(X.shape)
        i=0
        for X_t in X:
            #print(X_t.shape);print(hidden_state.shape);exit()
            if i>0: 
                hidden_state = torch.zeros(1,len(X),self.n_hidden) if len(hidden_s)<=i else hidden_s[i]
                cell_state = torch.zeros(1, len(X), self.n_hidden) if len(cell_s)<=i else cell_s[i]
            input_gate = torch.sigmoid(self.W_ii(X_t)+self.b_ii+self.W_hi(hidden_state)+self.b_hi)
            forget_gate = torch.sigmoid(self.W_if(X_t)+self.b_if+self.W_hf(hidden_state)+self.b_hf)
            cell_gate = torch.tanh(self.W_ig(X_t)+self.b_ig+self.W_hg(hidden_state)+self.b_hg)
            output_gate = torch.sigmoid(self.W_io(X_t)+self.b_io+self.W_ho(hidden_state)+self.b_ho)
            cell_state = forget_gate*cell_state+input_gate*cell_gate;
            hidden_state = output_gate*torch.tanh(cell_state);
            if len(cell_s)<=i:
                hidden_s.append(hidden_state.detach())
                cell_s.append(cell_state.detach())
            else:
                hidden_s[i]=hidden_state.detach()
                cell_s[i]=cell_state.detach()
            i=i+1
        outputs=output_gate[-1]
        model=self.W(outputs)+self.b;
        if len(tensors)<=self.numero+1: 
            tensors.append(hidden_s);
            cells.append(cell_s);
        else: 
            tensors[self.numero+1]=hidden_s;
            cells[self.numero+1]=cell_s;
        return torch.cat([hidden_state[0],model],dim=1);

class TextLSTM_3(nn.Module):
    def __init__(self,n_class,emb_size,n_hidden,numero:int):
        super(TextLSTM_3,self).__init__()
        self.numero=numero
        # self.C=nn.Embedding(n_class+n_hidden,embedding_dim=emb_size);
        self.W=nn.Linear(n_hidden,n_class,bias=False)
        self.b=nn.Parameter(torch.ones([n_class]))

        self.W_ii=nn.Linear(n_hidden+n_class,n_hidden,bias=False)
        self.W_if=nn.Linear(n_hidden+n_class,n_hidden,bias=False)
        self.W_ig=nn.Linear(n_hidden+n_class,n_hidden,bias=False)
        self.W_io=nn.Linear(n_hidden+n_class,n_hidden,bias=False)
        self.b_ii=nn.Parameter(torch.ones([n_hidden]))
        self.b_if=nn.Parameter(torch.ones([n_hidden]))
        self.b_ig=nn.Parameter(torch.ones([n_hidden]))
        self.b_io=nn.Parameter(torch.ones([n_hidden]))

        self.W_hi=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hf=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_hg=nn.Linear(n_hidden,n_hidden,bias=False)
        self.W_ho=nn.Linear(n_hidden,n_hidden,bias=False)
        self.b_hi=nn.Parameter(torch.ones([n_hidden]))
        self.b_hf=nn.Parameter(torch.ones([n_hidden]))
        self.b_hg=nn.Parameter(torch.ones([n_hidden]))
        self.b_ho=nn.Parameter(torch.ones([n_hidden]))
        self.n_hidden=n_hidden
    def forward(self,X):
        # X=self.C(X);
        hidden_s= [] if len(tensors)<=self.numero+1 else tensors[self.numero+1]
        cell_s=[] if len(cells)<=self.numero+1 else cells[self.numero+1]
        hidden_state = torch.zeros(1, len(X), self.n_hidden) if len(hidden_s)==0 else hidden_s[0]  # [num_layers(=1) * num_directions(=1), batch_size, n_hidden]
        cell_state = torch.zeros(1, len(X), self.n_hidden) if len(cell_s)==0 else cell_s[0]
        #X=X.transpose(0,1);
        i=0;
        for X_t in X:
            if i>0: 
                hidden_state = torch.zeros(1,len(X),self.n_hidden) if len(hidden_s)<=i else hidden_s[i]
                cell_state = torch.zeros(1, len(X), self.n_hidden) if len(cell_s)<=i else cell_s[i]
            input_gate = torch.sigmoid(self.W_ii(X_t)+self.b_ii+self.W_hi(hidden_state)+self.b_hi)
            forget_gate = torch.sigmoid(self.W_if(X_t)+self.b_if+self.W_hf(hidden_state)+self.b_hf)
            cell_gate = torch.tanh(self.W_ig(X_t)+self.b_ig+self.W_hg(hidden_state)+self.b_hg)
            output_gate = torch.sigmoid(self.W_io(X_t)+self.b_io+self.W_ho(hidden_state)+self.b_ho)
            cell_state = forget_gate*cell_state+input_gate*cell_gate;
            hidden_state = output_gate*torch.tanh(cell_state);
            if len(cell_s)<=i:
                hidden_s.append(hidden_state.detach())
                cell_s.append(cell_state.detach())
            else:
                hidden_s[i]=hidden_state.detach()
                cell_s[i]=cell_state.detach()
            i=i+1
        outputs=output_gate[-1]
        model=self.W(outputs)+self.b;
        if len(tensors)<=self.numero+1:
            tensors.append(hidden_s);
            cells.append(cell_s)
        else:
            tensors[-1]=hidden_s;
            cells[-1]=cell_s;
        return model;
def train_LSTMlm(n_class,emb_size,mid_layers=5):
    # model = TextLSTM(n_class,emb_size,n_hidden)
    model = nn.Sequential()
    model.add_module("first layer",TextLSTM_1(n_class,emb_size,n_hidden));
    for _ in range(mid_layers):
        model.add_module(str(_+1)+" middle layer",TextLSTM_2(n_class,emb_size,n_hidden,_))
    model.add_module("last layer",TextLSTM_3(n_class,emb_size,n_hidden,mid_layers))
    model.to(device)
    print(model)


    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learn_rate)
    
    # Training
    batch_number = len(all_input_batch)
    tensors.clear()
    cells.clear()
    for epoch in range(all_epoch):
        count_batch = 0
        for input_batch, target_batch in zip(all_input_batch, all_target_batch):
            optimizer.zero_grad()

            # input_batch : [batch_size, n_step, n_class]
            
            output = model(input_batch)
            

            # output : [batch_size, n_class], target_batch : [batch_size] (LongTensor, not one-hot)
            loss = criterion(output, target_batch)
            ppl = math.exp(loss.item())
            if (count_batch + 1) % 25 == 0:
                print('Epoch:', '%04d' % (epoch + 1), 'Batch:', '%02d' % (count_batch + 1), f'/{batch_number}',
                      'loss =', '{:.6f}'.format(loss), 'ppl =', '{:.6}'.format(ppl))

            loss.backward()
            optimizer.step()

            count_batch += 1

        print('Epoch:', '%04d' % (epoch + 1), 'Batch:', '%02d' % (count_batch + 1), f'/{batch_number}',
                'loss =', '{:.6f}'.format(loss), 'ppl =', '{:.6}'.format(ppl))

        # valid after training one epoch
        all_valid_batch, all_valid_target = give_valid_test.give_valid(data_root, word2number_dict, n_step)
        all_valid_batch = torch.LongTensor(all_valid_batch).to(device)  # list to tensor
        all_valid_target = torch.LongTensor(all_valid_target).to(device)

        total_valid = len(all_valid_target)*128  # valid and test batch size is 128
        with torch.no_grad():
            total_loss = 0
            count_loss = 0
            for valid_batch, valid_target in zip(all_valid_batch, all_valid_target):
                valid_output = model(valid_batch)
                valid_loss = criterion(valid_output, valid_target)
                total_loss += valid_loss.item()
                count_loss += 1
          
            print(f'Valid {total_valid} samples after epoch:', '%04d' % (epoch + 1), 'loss =',
                  '{:.6f}'.format(total_loss / count_loss),
                  'ppl =', '{:.6}'.format(math.exp(total_loss / count_loss)))

        if (epoch+1) % save_checkpoint_epoch == 0:
            torch.save(model, f'models/LSTMlm_model_epoch_{epoch+1}.ckpt')

def test_LSTMlm(select_model_path):
    model = torch.load(select_model_path, map_location="cpu")  #load the selected model
    model.to(device)

    #load the test data
    all_test_batch, all_test_target = give_valid_test.give_test(data_root, word2number_dict, n_step)
    all_test_batch = torch.LongTensor(all_test_batch).to(device)  # list to tensor
    all_test_target = torch.LongTensor(all_test_target).to(device)
    total_test = len(all_test_target)*128  # valid and test batch size is 128
    model.eval()
    criterion = nn.CrossEntropyLoss()
    total_loss = 0
    count_loss = 0
    for test_batch, test_target in zip(all_test_batch, all_test_target):
        test_output = model(test_batch)
        test_loss = criterion(test_output, test_target)
        total_loss += test_loss.item()
        count_loss += 1

    print(f"Test {total_test} samples with {select_model_path}……………………")
    print('loss =','{:.6f}'.format(total_loss / count_loss),
                  'ppl =', '{:.6}'.format(math.exp(total_loss / count_loss)))

if __name__ == '__main__':
    n_step = 5 # number of cells(= number of Step)
    n_hidden = 128 # number of hidden units in one cell
    batch_size = 128 # batch size
    learn_rate = 0.0005
    all_epoch = 5 #the all epoch for training
    emb_size = 256 #embeding size
    save_checkpoint_epoch = 5 # save a checkpoint per save_checkpoint_epoch epochs !!! Note the save path !!!
    data_root = 'penn_small'
    train_path = os.path.join(data_root, 'train.txt') # the path of train dataset

    print("print parameter ......")
    print("n_step:", n_step)
    print("n_hidden:", n_hidden)
    print("batch_size:", batch_size)
    print("learn_rate:", learn_rate)
    print("all_epoch:", all_epoch)
    print("emb_size:", emb_size)
    print("save_checkpoint_epoch:", save_checkpoint_epoch)
    print("train_data:", data_root)

    word2number_dict, number2word_dict = make_dict(train_path)
    #print(word2number_dict)

    print("The size of the dictionary is:", len(word2number_dict))
    n_class = len(word2number_dict)  #n_class (= dict size)

    print("generating train_batch ......")
    all_input_batch, all_target_batch = make_batch(train_path, word2number_dict, batch_size, n_step)  # make the batch
    train_batch_list = [all_input_batch, all_target_batch]
    
    print("The number of the train batch is:", len(all_input_batch))
    all_input_batch = torch.LongTensor(all_input_batch).to(device)   #list to tensor
    all_target_batch = torch.LongTensor(all_target_batch).to(device)
    # print(all_input_batch.shape)
    # print(all_target_batch.shape)
    all_input_batch = all_input_batch.reshape(-1, batch_size, n_step)
    all_target_batch = all_target_batch.reshape(-1, batch_size)

    print("\nTrain the LSTMLM……………………")
    train_LSTMlm(n_class,emb_size,1)

    print("\nTest the LSTMLM……………………")
    select_model_path = "models/LSTMlm_model_epoch_5.ckpt"
    test_LSTMlm(select_model_path)