models.py

import copy
import torch
from torch import nn
import numpy as np


class SummarisationModel(nn.Module):

    def __init__(self):

        super().__init__()

        self.cnns = [
            nn.Conv1d(200, 50, 1, padding="valid", groups=1),
            nn.Conv1d(200, 50, 2, padding="valid", groups=1),
            nn.Conv1d(200, 50, 3, padding="valid", groups=1),
            nn.Conv1d(200, 50, 4, padding="valid", groups=1),
            nn.Conv1d(200, 50, 5, padding="valid", groups=1),
            nn.Conv1d(200, 50, 6, padding="valid", groups=1),
            nn.Conv1d(200, 50, 7, padding="valid", groups=1)
            ]

        self.document_encoder = nn.LSTM(350, 600) 

        self.sentence_extractor = nn.LSTM(350, 600)

        self.projector = nn.Linear(600, 2)

        self.softmax = nn.Softmax(dim=-1)
    
    def encode_sentences(self, document):

        convolutions = []
        for cnn in self.cnns:
            convolutions.append(cnn(document.transpose(1,2)).amax(dim=2))

        return torch.cat(convolutions, dim=1)

    def encode_document(self, encoded_sentences):

        _, (hidden_state, cell_state) = self.document_encoder(encoded_sentences.flip(dims=(0,)))

        return hidden_state, cell_state

    def encode(self, document):

        encoded_sentences = self.encode_sentences(document)

        return encoded_sentences, self.encode_document(encoded_sentences)

    def forward(self, document, k=3):

        encoded_sentences, states = self.encode(document)

        logits = self.projector(self.sentence_extractor(encoded_sentences, states)[0])
        
        probs = self.softmax(logits)[:,0]
        
        if k < len(probs):
            return probs.topk(k).indices, probs # handle doc weniger als 3 sents? 
        return torch.arange(len(probs)), probs

    def test(self, dataset):



class ActorOnlySummarisationModel(SummarisationModel):
    
    def _train(self, dataset, epochs=20, batch_size=20, learning_rate=0.001, shuffle=True):
        
        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)

        training_dataloader = torch.utils.data.DataLoader(dataset.train, batch_size=batch_size, shuffle=shuffle)
        test_dataloader = torch.utils.data.DataLoader(dataset.test, batch_size=batch_size, shuffle=shuffle)

        since = time.time()
        val_rouge_history = []
    
        best_rouge = 0.0
        best_model_wts = copy.deepcopy(self.state_dict())
        
        for epoch in range(epochs):

            print('Epoch {}/{}'.format(epoch, epochs - 1))
            print('-' * 10)
            
            # training phase of the epoch
            running_loss = 0.0
            running_rouge = 0.0

            self.train()
            for batch in training_dataloader:

                optimizer.zero_grad()

                for datapoint in batch:
                    
                    top_indices, probs = self.__call__(datapoint.document)

                    o = datapoint.p_searchspace @ torch.log(probs) + datapoint.n_searchspace @ torch.log(1 - probs)

                    idx_sample = torch.argmax(o)

                    loss = - datapoint.top_rouge[idx_sample] * o[idx_sample]

                    loss.backward()

                    # training statistics (train dataset)
                    running_loss += loss.item()
                    running_rouge += datapoint.top_rouge[idx_sample] # größer als val/test rouge = indiz for searchspace funtionality

                optimizer.step()
            
            epoch_loss = running_loss/ len(training_dataloader.dataset)
            epoch_rouge = running_rouge/ len(training_dataloader.dataset) # abh. von __len__ of PreprocessedDataSet
            print('Train Loss: {:.4f} Rouge Score: {:.4f}'.format(epoch_loss, epoch_rouge))

            # validation phase of the epoch
            self.eval()
            running_rouge = 0.0
            with torch.no_grad():

                for datapoint in PreprocessedDataSet.validation:
                    
                    top_indices, probs = self.__call__(datapoint.document)
                    running_rouge += rouge(select_elements(datapoint.raw_document, top_indices), datapoint.raw_summary)  
                    # vgl. train rouge for searchspace funtionality
                epoch_rouge = running_rouge/ len(PreprocessedDataSet.validation)
                val_rouge_history.append(epoch_rouge)    
                print('Validation Rouge Score: {:.4f}'.format(epoch_rouge))
            
            # epoch completed, deep copy the best model sofar
            if epoch_rouge > best_rouge:
                best_rouge = epoch_rouge
                best_model_wts = copy.deepcopy(self.state_dict())

        # training completed
        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        print('Best val rouge: {:4f}'.format(best_rouge))
        # write val_rouge_history in file

        # load best model weights
        self.load_state_dict(best_model_wts)
        

class SummarisationModelWithCrossEntropyLoss(SummarisationModel):
    
    def _train(self, dataset, epochs=20, batch_size=20, learning_rate=0.001, shuffle=True):
        
        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
        loss_fn = nn.BCELoss(reduction='sum')
        
        for _ in range(epochs):
            
            training_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
            
            for batch in training_dataloader:

                optimizer.zero_grad()

                for datapoint in batch:
                    
                    _, probs = self.__call__(datapoint.document)

                    loss = loss_fn(probs, datapoint.bin_summary)

                    loss.backward()
                
                optimizer.step()

        # eval

class ActorCriticSummarisationModel(SummarisationModel):

    def __init__(self, )



class Critic(nn.Module):

    def __init__(self, model, steepness=8, denoise=100):

        super().__init__()

        self.steepness = steepness
        self.denoise = denoise

        model = copy.deepcopy(model)
        #model.eval()
        for param in model.parameters():
            param.requires_grad = False

        self.document_encoder = model.encode_document
 
        self.layer_1 = nn.Linear(1200, 600)
        self.layer_2 = nn.Linear(600, 600)
        self.layer_3 = nn.Linear(600, 1)

        W_1 = torch.cat((torch.eye(600), -torch.eye(600)), 1)
        W_2 = torch.eye(600)
        W_3 = torch.ones(600)

        self.layer_1.weight.data /= self.denoise
        self.layer_1.weight.data += W_1
        self.layer_2.weight.data /= self.denoise
        self.layer_2.weight.data += W_2
        self.layer_3.weight.data /= self.denoise
        self.layer_3.weight.data += W_3
        self.layer_3.weight.data /= 600

    def forward(self, encoded_sentences_1, encoded_sentences_2):

        _, document_vec_1 = self.document_encoder(encoded_sentences_1)
        _, document_vec_2 = self.document_encoder(encoded_sentences_2)

        double_document = torch.cat((torch.squeeze(document_vec_1), torch.squeeze(document_vec_2)), dim=-1)

        return torch.tanh(self.steepness*nn.functional.relu(self.layer_3(
            nn.functional.relu(self.layer_2(
            utils.gaussian(self.layer_1(double_document)))))))
    
    def _train(self, dataset, epochs=200, batch_size=20, learning_rate=0.001, shuffle=True, pos_samples=0.5):

        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
        loss_fn = nn.MSELoss()

        for _ in range(epochs):
            
            training_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)

            for batch in training_dataloader:

                optimizer.zero_grad()
            
                    for datapoint in batch:

                        r = np.random.random()
                        if r > pos_samples:

                            k = np.random.choice(len(datapoint.p_searchspace)) 
                            sample = datapoint.sent_vecs.masked_select(datapoint.p_searchspace[k].bool())  # not padded sent embeddngs

                            score = self.__call__(sample, datapoint.gold_sent_vecs) 
                            loss = loss_fn(score, datapoint.top_rouge[k])

                        else:
                            if len(datapoint.sent_vecs) >= 3: 
                                narray = np.random.choice(len(datapoint.sent_vecs), 3, replace = False) 
                                narray.sort()
                                sample = datapoint.sent_vecs[narray]
                            else:
                                continue # handle len(sent_vecs) < 3 

                            score = self.__call__(sample, datapoint.gold_sent_vecs)
                            loss = loss_fn(score, utils.rouge(raw_document[narray]), raw_summary)) 
                            # rouge score berechnen für negative sample => besser wäre externes berechnen und speichern?

                        loss.backward()
                
                optimizer.step()
            
            # eval

            # test with rouge