<jupyter_start><jupyter_text>Retrieve Poetry Poetry Retriever using the Poly-encoder Transformer architecture (Humeau et al., 2019) for retrieval<jupyter_code># This notebook is based on :
# https://aritter.github.io/CS-7650/
# This Project was developed at the Georgia Institute of Technology by <NAME> (<EMAIL>), 
# borrowing  from the Neural Machine Translation Project (Project 2) 
# of the UC Berkeley NLP course https://cal-cs288.github.io/sp20/
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import torch
from torch.jit import script, trace
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
import numpy as np
import csv
import random
import re
import os
import unicodedata
import codecs
from io import open
import itertools
import math
import pickle
import statistics
import sys
from functools import partial

from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
import tqdm
import nltk
#from google.colab import files
# General util functions
def make_dir_if_not_exists(directory):
	if not os.path.exists(directory):
		logging.info("Creating new directory: {}".format(directory))
		os.makedirs(directory)

def print_list(l, K=None):
	# If K is given then only print first K
	for i, e in enumerate(l):
		if i == K:
			break
		print(e)
	print()

def remove_multiple_spaces(string):
	return re.sub(r'\s+', ' ', string).strip()

def save_in_pickle(save_object, save_file):
	with open(save_file, "wb") as pickle_out:
		pickle.dump(save_object, pickle_out)

def load_from_pickle(pickle_file):
	with open(pickle_file, "rb") as pickle_in:
		return pickle.load(pickle_in)

def save_in_txt(list_of_strings, save_file):
	with open(save_file, "w") as writer:
		for line in list_of_strings:
			line = line.strip()
			writer.write(f"{line}\n")

def load_from_txt(txt_file):
	with open(txt_file, "r") as reader:
		all_lines = list()
		for line in reader:
			line = line.strip()
			all_lines.append(line)
		return all_lines
import pandas as pd

print(torch.cuda.is_available())
if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
print("Using device:", device)
bert_model_name = 'distilbert-base-uncased' 
# Bert Imports
from transformers import DistilBertTokenizer, DistilBertModel
#bert_model = DistilBertModel.from_pretrained(bert_model_name)
tokenizer = DistilBertTokenizer.from_pretrained(bert_model_name)<jupyter_output><empty_output><jupyter_text>Load Data  Poetry Database<jupyter_code>data_file = '../data/with_epoque.csv'
data = pd.read_csv(data_file)
print(len(data))
print(data.head())<jupyter_output>573
                                    author  \
0                      <NAME>   
1  DUCHESS OF NEW<NAME>   
2                           <NAME>   
3                           EDMUND SPENSER   
4                        <NAME>   

                                             content  \
0  Let the bird of loudest lay\nOn the sole Arabi...   
1  Sir Charles into my chamber coming in,\nWhen I...   
2  Our vice runs beyond all that old men saw,\nAn...   
3  Lo I the man, whose Muse whilome did maske,\nA...   
4  Long have I longd to see my love againe,\nStil...   

                                 poem name          age                  type  
0               The Phoenix and the Turtle  Renaissance  Mythology & Folklore  
1                 An Epilogue to the Above  Renaissance  Mythology & Folklore  
2                       Book 7, Epigram 42  Renaissance  Mythology & Folklore  
3  from The Faerie Queene: Book I, Canto I  Renaissance  Mytho[...]<jupyter_text>Dataset Preparation<jupyter_code>def make_data_training(df, char_max_line = 20):
    inputs = []
    context = []
    targets = []
    for i,rows in df.iterrows():
        splitted = rows['content'].split('\r\n')
        for line in splitted:
            if len(line.strip()) > 0 and len(line.split(' ')) <= char_max_line:
                inputs.append(line)
                targets.append(line)
                context.append(' '.join([str(rows['poem name'])]))
        
    return pd.DataFrame(list(zip(inputs, context, targets)),columns =['text', 'context','target'])


#Defining torch dataset class for poems
class PoemDataset(Dataset):
    def __init__(self, df):
        self.df = df

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        return self.df.iloc[idx]
df = make_data_training(data, char_max_line = 30)
pad_word = "<pad>"
bos_word = "<bos>"
eos_word = "<eos>"
unk_word = "<unk>"
sep_word = "sep"

pad_id = 0
bos_id = 1
eos_id = 2
unk_id = 3
sep_id = 4
    
def normalize_sentence(s):
    s = re.sub(r"([.!?])", r" \1", s)
    s = re.sub(r"[^a-zA-Z.!?]+", r" ", s)
    s = re.sub(r"\s+", r" ", s).strip()
    return s

class Vocabulary:
    def __init__(self):
        self.word_to_id = {pad_word: pad_id, bos_word: bos_id, eos_word:eos_id, unk_word: unk_id, sep_word: sep_id}
        self.word_count = {}
        self.id_to_word = {pad_id: pad_word, bos_id: bos_word, eos_id: eos_word, unk_id: unk_word, sep_id: sep_word}
        self.num_words = 5
    
    def get_ids_from_sentence(self, sentence):
        sentence = normalize_sentence(sentence)
        sent_ids = [bos_id] + [self.word_to_id[word.lower()] if word.lower() in self.word_to_id \
                               else unk_id for word in sentence.split()] + \
                               [eos_id]
        return sent_ids
    
    def tokenized_sentence(self, sentence):
        sent_ids = self.get_ids_from_sentence(sentence)
        return [self.id_to_word[word_id] for word_id in sent_ids]

    def decode_sentence_from_ids(self, sent_ids):
        words = list()
        for i, word_id in enumerate(sent_ids):
            if word_id in [bos_id, eos_id, pad_id]:
                # Skip these words
                continue
            else:
                words.append(self.id_to_word[word_id])
        return ' '.join(words)

    def add_words_from_sentence(self, sentence):
        sentence = normalize_sentence(sentence)
        for word in sentence.split():
            if word not in self.word_to_id:
                # add this word to the vocabulary
                self.word_to_id[word] = self.num_words
                self.id_to_word[self.num_words] = word
                self.word_count[word] = 1
                self.num_words += 1
            else:
                # update the word count
                self.word_count[word] += 1

vocab = Vocabulary()
for src in df['text']:
    vocab.add_words_from_sentence(src.lower())

print(f"Total words in the vocabulary = {vocab.num_words}")
class Poem_dataset(Dataset):
    """Single-Turn version of Cornell Movie Dialog Cropus dataset."""

    def __init__(self, poems, context,vocab, device):
        """
        Args:
            conversations: list of tuple (src_string, tgt_string) 
                         - src_string: String of the source sentence
                         - tgt_string: String of the target sentence
            vocab: Vocabulary object that contains the mapping of 
                    words to indices
            device: cpu or cuda
        """
        l = []
        
        for i in range(len(poems)):
            l.append( ( context[i] + ' sep ' + poems[i] , poems[i] ))
        
        self.conversations = l.copy()
        self.vocab = vocab
        self.device = device

        def encode(src, tgt):
            src_ids = self.vocab.get_ids_from_sentence(src)
            tgt_ids = self.vocab.get_ids_from_sentence(tgt)
            return (src_ids, tgt_ids)

        # We will pre-tokenize the conversations and save in id lists for later use
        self.tokenized_conversations = [encode(src, tgt) for src, tgt in self.conversations]
        
    def __len__(self):
        return len(self.conversations)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()

        return {"conv_ids":self.tokenized_conversations[idx], "conv":self.conversations[idx]}

def collate_fn(data):
    """Creates mini-batch tensors from the list of tuples (src_seq, tgt_seq).
    We should build a custom collate_fn rather than using default collate_fn,
    because merging sequences (including padding) is not supported in default.
    Seqeuences are padded to the maximum length of mini-batch sequences (dynamic padding).
    Args:
        data: list of dicts {"conv_ids":(src_ids, tgt_ids), "conv":(src_str, trg_str)}.
            - src_ids: list of src piece ids; variable length.
            - tgt_ids: list of tgt piece ids; variable length.
            - src_str: String of src
            - tgt_str: String of tgt
    Returns: dict { "conv_ids":     (src_ids, tgt_ids), 
                    "conv":         (src_str, tgt_str), 
                    "conv_tensors": (src_seqs, tgt_seqs)}
            src_seqs: torch tensor of shape (src_padded_length, batch_size).
            tgt_seqs: torch tensor of shape (tgt_padded_length, batch_size).
            src_padded_length = length of the longest src sequence from src_ids
            tgt_padded_length = length of the longest tgt sequence from tgt_ids
    """
    # Sort conv_ids based on decreasing order of the src_lengths.
    # This is required for efficient GPU computations.
    src_ids = [torch.LongTensor(e["conv_ids"][0]) for e in data]
    tgt_ids = [torch.LongTensor(e["conv_ids"][1]) for e in data]
    src_str = [e["conv"][0] for e in data]
    tgt_str = [e["conv"][1] for e in data]
    data = list(zip(src_ids, tgt_ids, src_str, tgt_str))
    data.sort(key=lambda x: len(x[0]), reverse=True)
    src_ids, tgt_ids, src_str, tgt_str = zip(*data)


    # Pad the src_ids and tgt_ids using token pad_id to create src_seqs and tgt_seqs
    
    # Implementation tip: You can use the nn.utils.rnn.pad_sequence utility
    # function to combine a list of variable-length sequences with padding.
    
    # YOUR CODE HERE
    src_seqs = nn.utils.rnn.pad_sequence(src_ids, padding_value = pad_id,
                                         batch_first = False)
    tgt_seqs = nn.utils.rnn.pad_sequence(tgt_ids, padding_value = pad_id, 
                                         batch_first = False)
    
    src_padded_length = len(src_seqs[0])
    tgt_padded_length = len(tgt_seqs[0])
    return {"conv_ids":(src_ids, tgt_ids), "conv":(src_str, tgt_str), "conv_tensors":(src_seqs.to(device), tgt_seqs.to(device))}
# Create the DataLoader for all_conversations

all_poems = df['text'].tolist()
context = df['context'].tolist()

dataset = Poem_dataset(all_poems, context, vocab, device)

batch_size = 5

data_loader = DataLoader(dataset=dataset, batch_size=batch_size, 
                               shuffle=True, collate_fn=collate_fn)

for src, tgt in dataset.conversations[:3]:
    sentence = src
    word_tokens = vocab.tokenized_sentence(sentence)
    # Automatically adds bos_id and eos_id before and after sentence ids respectively
    word_ids = vocab.get_ids_from_sentence(sentence)
    print(sentence)
    print(word_tokens)
    print(word_ids)
    print(vocab.decode_sentence_from_ids(word_ids))
    print()

word = "the"
word_id = vocab.word_to_id[word.lower()]
print(f"Word = {word}")
print(f"Word ID = {word_id}")
print(f"Word decoded from ID = {vocab.decode_sentence_from_ids([word_id])}")
# Test one batch of training data
first_batch = next(iter(data_loader))
print(f"Testing first training batch of size {len(first_batch['conv'][0])}")
print(f"List of source strings:")
print_list(first_batch["conv"][0])
print(f"Tokenized source ids:")
print_list(first_batch["conv_ids"][0])
print(f"Padded source ids as tensor (shape {first_batch['conv_tensors'][0].size()}):")
print(first_batch["conv_tensors"][0])
def transformer_collate_fn(batch, tokenizer):
    bert_vocab = tokenizer.get_vocab()
    bert_pad_token = bert_vocab['[PAD]']
    bert_unk_token = bert_vocab['[UNK]']
    bert_cls_token = bert_vocab['[CLS]']
    inputs, masks_input, outputs, masks_output = [], [], [], []

    sentences, masks_sentences, targets, masks_targets = [], [], [], []
    for data in batch:

        tokenizer_output = tokenizer([data['text']])
        tokenized_sent = tokenizer_output['input_ids'][0]
        
        tokenizer_target = tokenizer([data['target']])
        tokenized_sent_target = tokenizer_target['input_ids'][0]
        
        mask_sentence = tokenizer_output['attention_mask'][0]
        mask_target = tokenizer_target['attention_mask'][0]
        sentences.append(torch.tensor(tokenized_sent))
        targets.append(torch.tensor(tokenized_sent_target))
        masks_targets.append(torch.tensor(mask_targets))
        masks_sentences.append(torch.tensor(mask_sentences))
    sentences = pad_sequence(sentences, batch_first=True, padding_value=bert_pad_token)
    targets = pad_sequence(targets, batch_first=True, padding_value=bert_pad_token)
    masks = pad_sequence(masks, batch_first=True, padding_value=0.0)
    return sentences, targets, masks
#create pytorch dataloaders from train_dataset, val_dataset, and test_datset
batch_size=5
train_dataloader = DataLoader(dataset,batch_size=batch_size,collate_fn=partial(transformer_collate_fn, tokenizer=tokenizer), shuffle = True)
#tokenizer.batch_decode(transformer_collate_fn(train_dataset,tokenizer)[0], skip_special_tokens=True)<jupyter_output><empty_output><jupyter_text>Polyencoder Model<jupyter_code>#torch.cuda.empty_cache()
#bert1 = DistilBertModel.from_pretrained(bert_model_name)
#bert2 = DistilBertModel.from_pretrained(bert_model_name)

bert = DistilBertModel.from_pretrained(bert_model_name)
#Double Bert
class RetrieverPolyencoder(nn.Module):
    def __init__(self, contextBert, candidateBert, vocab, max_len = 300, hidden_dim = 768, out_dim = 64, num_layers = 2, dropout=0.1, device=device):
        super().__init__()

        self.device = device
        self.hidden_dim = hidden_dim
        self.max_len = max_len
        self.out_dim = out_dim
        
        # Context layers
        self.contextBert = contextBert
        self.contextDropout = nn.Dropout(dropout)
        self.contextFc = nn.Linear(self.hidden_dim, self.out_dim)
        
        # Candidates layers
        self.candidatesBert = candidateBert
        self.pos_emb = nn.Embedding(self.max_len, self.hidden_dim)
        self.candidatesDropout = nn.Dropout(dropout)
        self.candidatesFc = nn.Linear(self.hidden_dim, self.out_dim)
        
        self.att_dropout = nn.Dropout(dropout)


    def attention(self, q, k, v, vMask=None):
        w = torch.matmul(q, k.transpose(-1, -2))
        if vMask is not None:
            w *= vMask.unsqueeze(1)
            w = F.softmax(w, -1)
        w = self.att_dropout(w)
        score = torch.matmul(w, v)
        return score

    def score(self, context, context_mask, responses, responses_mask):
        """Run the model on the source and compute the loss on the target.

        Args:
            source: An integer tensor with shape (max_source_sequence_length,
                batch_size) containing subword indices for the source sentences.
            target: An integer tensor with shape (max_target_sequence_length,
                batch_size) containing subword indices for the target sentences.

        Returns:
            A scalar float tensor representing cross-entropy loss on the current batch
            divided by the number of target tokens in the batch.
            Many of the target tokens will be pad tokens. You should mask the loss 
            from these tokens using appropriate mask on the target tokens loss.
        """
        batch_size, nb_cand, seq_len = responses.shape
        # Context
        context_encoded = self.contextBert(context,context_mask)[-1]
        pos_emb = self.pos_emb(torch.arange(self.max_len).to(self.device))
        context_att = self.attention(pos_emb, context_encoded, context_encoded, context_mask)

        # Response
        responses_encoded = self.candidatesBert(responses.view(-1,responses.shape[2]), responses_mask.view(-1,responses.shape[2]))[-1][:,0,:]
        responses_encoded = responses_encoded.view(batch_size,nb_cand,-1)
        
        context_emb = self.attention(responses_encoded, context_att, context_att).squeeze() 
        dot_product = (context_emb*responses_encoded).sum(-1)
        
        return dot_product

    
    def compute_loss(self, context, context_mask, response, response_mask):
        """Run the model on the source and compute the loss on the target.

        Args:
            source: An integer tensor with shape (max_source_sequence_length,
                batch_size) containing subword indices for the source sentences.
            target: An integer tensor with shape (max_target_sequence_length,
                batch_size) containing subword indices for the target sentences.

        Returns:
            A scalar float tensor representing cross-entropy loss on the current batch
            divided by the number of target tokens in the batch.
            Many of the target tokens will be pad tokens. You should mask the loss 
            from these tokens using appropriate mask on the target tokens loss.
        """
        batch_size = context.shape[0]
        
        # Context
        context_encoded = self.contextBert(context,context_mask)[-1]
        pos_emb = self.pos_emb(torch.arange(self.max_len).to(self.device))
        context_att = self.attention(pos_emb, context_encoded, context_encoded, context_mask)

        # Response
        response_encoded = self.candidatesBert(response, response_mask)[-1][:,0,:]
        
        response_encoded = response_encoded.unsqueeze(0).expand(batch_size, batch_size, response_encoded.shape[1]) 
        context_emb = self.attention(response_encoded, context_att, context_att).squeeze() 
        dot_product = (context_emb*response_encoded).sum(-1)
        mask = torch.eye(batch_size).to(self.device)
        loss = F.log_softmax(dot_product, dim=-1) * mask
        loss = (-loss.sum(dim=1)).mean()
        return loss
#Single Bert
class RetrieverPolyencoder_single(nn.Module):
    def __init__(self, bert, max_len = 300, hidden_dim = 768, out_dim = 64, num_layers = 2, dropout=0.1, device=device):
        super().__init__()

        self.device = device
        self.hidden_dim = hidden_dim
        self.max_len = max_len
        self.out_dim = out_dim
        self.bert = bert
        
        # Context layers
        self.contextDropout = nn.Dropout(dropout)
        
        # Candidates layers
        self.pos_emb = nn.Embedding(self.max_len, self.hidden_dim)
        self.candidatesDropout = nn.Dropout(dropout)
        
        self.att_dropout = nn.Dropout(dropout)


    def attention(self, q, k, v, vMask=None):
        w = torch.matmul(q, k.transpose(-1, -2))
        if vMask is not None:
            w *= vMask.unsqueeze(1)
            w = F.softmax(w, -1)
        w = self.att_dropout(w)
        score = torch.matmul(w, v)
        return score

    def score(self, context, context_mask, responses, responses_mask):
        """Run the model on the source and compute the loss on the target.

        Args:
            source: An integer tensor with shape (max_source_sequence_length,
                batch_size) containing subword indices for the source sentences.
            target: An integer tensor with shape (max_target_sequence_length,
                batch_size) containing subword indices for the target sentences.

        Returns:
            A scalar float tensor representing cross-entropy loss on the current batch
            divided by the number of target tokens in the batch.
            Many of the target tokens will be pad tokens. You should mask the loss 
            from these tokens using appropriate mask on the target tokens loss.
        """
        batch_size, nb_cand, seq_len = responses.shape
        # Context
        context_encoded = self.bert(context,context_mask)[0][:,0,:]
        pos_emb = self.pos_emb(torch.arange(self.max_len).to(self.device))
        context_att = self.attention(pos_emb, context_encoded, context_encoded, context_mask)

        # Response
        responses_encoded = self.bert(responses.view(-1,responses.shape[2]), responses_mask.view(-1,responses.shape[2]))[0][:,0,:]
        responses_encoded = responses_encoded.view(batch_size,nb_cand,-1)
        response_encoded = self.candidatesFc(response_encoded)
        
        context_emb = self.attention(responses_encoded, context_att, context_att).squeeze() 
        dot_product = (context_emb*responses_encoded).sum(-1)
        
        return dot_product

    
    def compute_loss(self, context, context_mask, response, response_mask):
        """Run the model on the source and compute the loss on the target.

        Args:
            source: An integer tensor with shape (max_source_sequence_length,
                batch_size) containing subword indices for the source sentences.
            target: An integer tensor with shape (max_target_sequence_length,
                batch_size) containing subword indices for the target sentences.

        Returns:
            A scalar float tensor representing cross-entropy loss on the current batch
            divided by the number of target tokens in the batch.
            Many of the target tokens will be pad tokens. You should mask the loss 
            from these tokens using appropriate mask on the target tokens loss.
        """
        batch_size = context.shape[0]
        seq_len = response.shape[1]
        
        # Context
        context_encoded = self.bert(context,context_mask)[0][:,0,:]
        pos_emb = self.pos_emb(torch.arange(self.max_len).to(self.device))
        context_att = self.attention(pos_emb, context_encoded, context_encoded, context_mask)

        # Response
        print(response.shape)
        response_encoded = self.bert(response, response_mask)[0][:,0,:]
        print(response_encoded.shape)
        response_encoded = response_encoded.view(batch_size, -1)

        
        response_encoded = response_encoded.unsqueeze(0).expand(batch_size, batch_size, response_encoded.shape[1]) 
        context_emb = self.attention(response_encoded, context_att, context_att).squeeze() 
        dot_product = (context_emb*response_encoded).sum(-1)
        mask = torch.eye(batch_size).to(self.device)
        loss = F.log_softmax(dot_product, dim=-1) * mask
        loss = (-loss.sum(dim=1)).mean()
        return loss
#Bi-encoder
class RetrieverBiencoder(nn.Module):
    def __init__(self, bert):
        super().__init__()
        self.bert = bert
        
    def score(self, context, context_mask, responses, responses_mask):

        context_vec = self.bert(context, context_mask)[0][:,0,:]  # [bs,dim]

        batch_size, res_length = response.shape

        responses_vec = self.bert(responses_input_ids, responses_input_masks)[0][:,0,:]  # [bs,dim]
        responses_vec = responses_vec.view(batch_size, 1, -1)

        responses_vec = responses_vec.squeeze(1)        
        context_vec = context_vec.unsqueeze(1)
        dot_product = torch.matmul(context_vec, responses_vec.permute(0, 2, 1)).squeeze()
        return dot_product
    
    def compute_loss(self, context, context_mask, response, response_mask):

        context_vec = self.bert(context, context_mask)[0]  # [bs,dim]

        batch_size, res_length = response.shape

        responses_vec = self.bert(response, response_mask)[0][:,0,:]  # [bs,dim]
        #responses_vec = responses_vec.view(batch_size, 1, -1)
        
        print(context_vec.shape)
        print(responses_vec.shape)

        responses_vec = responses_vec.squeeze(1)
        dot_product = torch.matmul(context_vec, responses_vec.t())  # [bs, bs]
        mask = torch.eye(context.size(0)).to(context_mask.device)
        loss = F.log_softmax(dot_product, dim=-1) * mask
        loss = (-loss.sum(dim=1)).mean()
        return loss

def train(model, data_loader, num_epochs, model_file, learning_rate=0.0001):
    """Train the model for given µnumber of epochs and save the trained model in 
    the final model_file.
    """

    decoder_learning_ratio = 5.0
    #encoder_parameter_names = ['word_embedding', 'encoder']
    encoder_parameter_names = ['encode_emb', 'encode_gru', 'l1', 'l2']
                           
    encoder_named_params = list(filter(lambda kv: any(key in kv[0] for key in encoder_parameter_names), model.named_parameters()))
    decoder_named_params = list(filter(lambda kv: not any(key in kv[0] for key in encoder_parameter_names), model.named_parameters()))
    encoder_params = [e[1] for e in encoder_named_params]
    decoder_params = [e[1] for e in decoder_named_params]
    optimizer = torch.optim.AdamW([{'params': encoder_params},
                {'params': decoder_params, 'lr': learning_rate * decoder_learning_ratio}], lr=learning_rate)
    
    clip = 50.0
    for epoch in tqdm.notebook.trange(num_epochs, desc="training", unit="epoch"):
        # print(f"Total training instances = {len(train_dataset)}")
        # print(f"train_data_loader = {len(train_data_loader)} {1180 > len(train_data_loader)/20}")
        with tqdm.notebook.tqdm(
                data_loader,
                desc="epoch {}".format(epoch + 1),
                unit="batch",
                total=len(data_loader)) as batch_iterator:
            model.train()
            total_loss = 0.0
            for i, batch_data in enumerate(batch_iterator, start=1):
                source, mask_source, target, mask_target = batch_data["conv_tensors"]
                optimizer.zero_grad()
                loss = model.compute_loss(source, mask_source, target, mask_target)
                total_loss += loss.item()
                loss.backward()
                # Gradient clipping before taking the step
                _ = nn.utils.clip_grad_norm_(model.parameters(), clip)
                optimizer.step()

                batch_iterator.set_postfix(mean_loss=total_loss / i, current_loss=loss.item())
    # Save the model after training         
    torch.save(model.state_dict(), model_file)
# You are welcome to adjust these parameters based on your model implementation.
num_epochs = 10
batch_size = 32
learning_rate = 0.001
# Reloading the data_loader to increase batch_size

baseline_model = RetrieverBiencoder(bert).to(device)
train(baseline_model, train_dataloader, num_epochs, "baseline_model.pt",learning_rate=learning_rate)
# Download the trained model to local for future use
#files.download('baseline_model.pt')
baseline_model = RetrieverPolyencoder(bert1,bert2,vocab).to(device)
baseline_model.load_state_dict(torch.load("baseline_model3.pt", map_location=device))
vals = transformer_collate_fn(all_conversations[0:100],tokenizer)
i=3
scores = baseline_model.score(vals[0][i].unsqueeze(0).cuda(),vals[1][i].unsqueeze(0).cuda(),vals[2].unsqueeze(0).cuda(),vals[3].unsqueeze(0).cuda()).detach().cpu().numpy()
all_conversations[i][0]
all_conversations[np.argmax(scores)][1]
max_v = 100
vals = transformer_collate_fn(all_conversations[0:max_v],tokenizer)
correct = 0
for i in range(max_v):
    scores = baseline_model.score(vals[0][i].unsqueeze(0).cuda(),vals[1][i].unsqueeze(0).cuda(),vals[2].unsqueeze(0).cuda(),vals[3].unsqueeze(0).cuda()).detach().cpu().numpy()
    if np.argmax(scores)==i:
        correct+=1
    print(all_conversations[i][0])
    print(all_conversations[np.argmax(scores)][1]+"\n")
print(correct/max_v)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.102133536873478, 'fraction_numerical': 0.008396030967179152, 'mean_word_length': 2.9294487289345903, 'pattern_counts': {'":': 10, '<': 27, '<?xml version=': 0, '>': 28, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 11, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '10880', 'license': 'MIT', 'n_tokens_mistral': 8820, 'n_tokens_neox': 8522, 'n_words': 2079, 'path': 'notebooks/retrieve_poetry.ipynb', 'repo_name': 'M-Puig/Erato', 'size': '27595'}