train_cross_encoder.py

import torch

from optimizers import standard_optimizer
from tqdm import tqdm, trange
from parameters import RankingParser
from models.CrossEncoder import CrossEncoder
import random
import os
import pickle
from optimizers import noise
import logging
from transformers import AutoTokenizer
from torch.utils.data import DataLoader
logging.disable(logging.WARNING)
device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu")

max_candsize = 5
add_gold_mention = True


class Trainer:
    def __init__(self, params,  device):
        self.grad_acc_steps = params["gradient_accumulation_steps"]
        self.params = params
        self.evaluate_after = params["eval_interval"]
        #self.candidate_size = candidate_size
        self.device = device
        self.model = CrossEncoder(params,device)
        #self.tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-base-v2')
        self.tokenizer = AutoTokenizer.from_pretrained(params["base_model"])
        #self.collator = E5collator(tokenizer=self.tokenizer,queries=handler.queries,device=self.device)
        self.model.to(device)
        self.ms_marco_documents=None
        if self.params["dataset"]=="msmarco":
            gold_documents = pickle.load(open("data/msmarco/gold_documents", "rb"))
            self.ms_marco_documents = pickle.load(open("data/msmarco/eval_documents_100000", "rb"))
            self.ms_marco_documents.update(gold_documents)
        # self.optimizer,self.scheduler=self.getOptimizerAndSheduler()

    '''
    def get_train_test_split(self, test_split, samples):
        split = len(samples) * test_split
        train = samples[0:len(samples) - int(split)]
        test = samples[len(samples) - int(split):-1]
        return train, test

    '''
    '''
    def getDataLoaders(self, data):
        random.shuffle(data)
        train, test = self.get_train_test_split(0.05, data)
        test_dataloader = DataLoader(test, shuffle=True, batch_size=self.params["eval_batch_size"],
                                        collate_fn=self.collator.collate)
        train_dataloader = DataLoader(train, shuffle=True, batch_size=self.params["train_batch_size"],
                                        collate_fn=self.collator.collate)
        return train_dataloader, test_dataloader
    '''
    def getOptimizerAndSheduler(self, len_train_Data):
        optimizer = standard_optimizer.get_bert_optimizer([self.model], self.params["type_optimization"],
            self.params["learning_rate"],
                fp16=self.params.get("fp16"))
        scheduler = standard_optimizer.get_scheduler(self.params, optimizer, len_train_Data)
        return optimizer, scheduler
    def collate(self, batch,  num_candidates,num_noise_labels):
        samples=[]
        labels=[]
        for doc in batch:
            if len(doc["correct_repr"])>0:
                query=doc["query"]
                correct=random.choice(doc["correct_repr"])
                candidates=random.sample(doc["candidate_repr"],num_candidates-1)
                candidates.append(correct)
                random.shuffle(candidates)
                correct_ind=candidates.index(correct)
                for i in range(0,len(candidates)):
                    samples.append("query: "+query+"[SEP]candidate: "+candidates[i])
                labels.append(correct_ind)
        for i in range(num_noise_labels):
            ind_to_update=random.randrange(len(batch))
            update_ind = labels[ind_to_update]
            while update_ind == labels[ind_to_update]:
                update_ind = random.randrange(num_candidates)
            labels[ind_to_update] = update_ind
        return samples,labels
    def collate_ms_marco(self, batch,  num_candidates,num_noise_labels):
        samples=[]
        labels=[]
        for doc in batch:
            if len(doc["correct"])>0:
                query=doc["query"]
                correct=random.choice(doc["correct"])
                candidates=random.sample(doc["candidates"],num_candidates-1)
                candidates.append(correct)
                candidates_repr=["passage: title: " + self.ms_marco_documents[cand][1] + "[SEP] context: " + self.ms_marco_documents[cand][2]
                            for cand in candidates]
                random.shuffle(candidates_repr)
                correct_ind=candidates.index(correct)
                for i in range(0,len(candidates)):
                    samples.append("query: "+query+"[SEP]candidate: "+candidates_repr[i])
                labels.append(correct_ind)
        for i in range(num_noise_labels):
            ind_to_update=random.randrange(len(batch))
            update_ind = labels[ind_to_update]
            while update_ind == labels[ind_to_update]:
                update_ind = random.randrange(num_candidates)
            labels[ind_to_update] = update_ind
        return samples,labels

    def make_forward_pass(self, batch):
        if self.params["dataset"] == "msmarco":
            batch,labels=self.collate_ms_marco(batch,num_candidates=15, num_noise_labels=0)
        else:
            batch,labels=self.collate(batch,num_candidates=15, num_noise_labels=0)
        #documents = batch[1]
        batch=self.tokenizer(batch, max_length=512, padding=True, truncation=True, return_tensors='pt')
        batch.to(self.device)
        labels=torch.tensor(labels,device=self.device)
        loss, logits = self.model(batch,15,labels)
        return logits, loss
    def evaluate(self,eval_data,num_candidates,k):
        self.model.eval()
        all_found=0
        all_not_found=0
        all_rr = []
        for sample in tqdm(eval_data):
            query = "query: "+sample["query"]
            if self.params["dataset"] == "msmarco":
                correct = sample["correct"]
                correct=["passage: title: " + self.ms_marco_documents[cand][1] + "[SEP] context: " +
                                   self.ms_marco_documents[cand][2]
                                   for cand in correct]
                #candidates = correct.copy()
                candidates = random.sample(sample["candidates"], num_candidates - len(correct))
                candidates = ["passage: title: " + self.ms_marco_documents[cand][1] + "[SEP] context: " +
                                   self.ms_marco_documents[cand][2]
                                   for cand in candidates]
                candidates.extend(correct)
                batch = [query + "[SEP]candidate: " + el for el in candidates]
                batch = self.tokenizer(batch, max_length=512, padding=True, truncation=True, return_tensors='pt')
            else:
                correct = sample["correct_repr"]
                candidates = correct.copy()
                candidates.extend(random.sample(sample["candidate_repr"], num_candidates - len(candidates)))
                batch=[query+"[SEP]candidate: "+el for el in candidates]
                batch = self.tokenizer(batch, max_length=512, padding=True, truncation=True, return_tensors='pt')
            batch.to(self.device)
            scores = self.model(batch)[1]
            best=torch.topk(scores,k).indices.tolist()
            prediction=[candidates[el] for el in best]

            for en in correct:
                if en in prediction:
                    all_found += 1
                else:
                    all_not_found += 1
            reciprocal_rank = 0
            for rank, pred in enumerate(prediction, start=1):
                if pred in correct:
                    reciprocal_rank = 1 / rank
                    break

            all_rr.append(reciprocal_rank)

        mrr = sum(all_rr) / len(all_rr) if all_rr else 0
        results = all_found / (all_not_found + all_found)

        return results, mrr




def load_training():
    parser = RankingParser(add_model_args=True)
    parser.add_training_args()
    parser.add_eval_args()

    # args = argparse.Namespace(**params)
    args = parser.parse_args()
    print(args)
    params = args.__dict__

    global device
    device = torch.device(
        "cuda:" + str(params["gpu_id"]) if torch.cuda.is_available() else "cpu")
    # for lc-quad
    if params["dataset"] == "lcquad":
        train_data = pickle.load(open("data/cross_encoder_data/lcquad/lcquad_train_cross_encoder_data.pkl", "rb"))
        test_data = pickle.load(open("data/cross_encoder_data/lcquad/lcquad_test_cross_encoder_data.pkl", "rb"))
    # for mintaka
    elif params["dataset"] == "mintaka":
        train_data = pickle.load(open("data/cross_encoder_data/mintaka/train_cross_encoder_data.pkl", "rb"))
        test_data = pickle.load(open("data/cross_encoder_data/mintaka/test_cross_encoder_data.pkl", "rb"))
    elif params["dataset"] == "aida":
        train_data = pickle.load(open("data/cross_encoder_data/aida/train_crossencoder.pkl", "rb"))
        test_data = pickle.load(open("data/cross_encoder_data/aida/test_crossencoder.pkl", "rb"))
    elif params["dataset"] == "msmarco":
        train_data = pickle.load(open("data/cross_encoder_data/msmarco/train_cross_encoder_data.pkl", "rb"))
        test_data = pickle.load(open("data/cross_encoder_data/msmarco/eval_cross_encoder_data.pkl", "rb"))

    # for E5
    trainer=Trainer(params,device)

    #train_dataloader = DataLoader(list(entities.keys()), shuffle=True, batch_size=1,
    #                              )
    # train_inst=Trainer.TrainerRanker(params=params, evaluate_after_batch=params["eval_interval"], device=device)

    optimizer, scheduler = trainer.getOptimizerAndSheduler(10000)


    return trainer,  optimizer, scheduler,train_data, test_data


def save_model(model, tokenizer, output_dir):
    """Saves the model and the tokenizer used in the output directory."""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    model_to_save = model.module if hasattr(model, "module") else model
    output_model_file = os.path.join(output_dir, "pytorch_model.bin")
    output_config_file = os.path.join(output_dir, "model_config")
    torch.save(model_to_save.state_dict(), output_model_file)
    # model_to_save.config.to_json_file(output_config_file)
    tokenizer.save_vocabulary(output_dir)



def collate(batch):
    return batch

def train(epochs):
    # trainer,evaluator, train_dataloader, optimizer, scheduler = load_train_only_Graph_Model(device)
    trainer,  optimizer, scheduler,train_data,test_data = load_training()

    # print(evaluator.evaluate(trainer.model))
    #results, mrr = trainer.evaluate(test_data,30,10)
    results, mrr = trainer.evaluate(test_data, 30, 10)
    trainer.model.train()
    # results, mrr = evaluator.evaluate_mrr(trainer.model)
    #print(results)
    #print(mrr)
    dataloader=DataLoader(train_data, shuffle=True, batch_size=10,
               collate_fn=collate)
    iter_ = tqdm(dataloader, desc="Training")
    for e in range(epochs):
        num_batch = 0
        # step=0
        for step, batch in enumerate(iter_):
            # batch=data_processing.create_batch_ent(batch[0],list(entities[batch[0]]),random.sample(list(documents),1000),doc_to_ent)
            # batch = data_processing.create_batch_index_document(batch[0], entities,  encoding_map,
            #                                           index, doc_to_ent)
            logits, loss = trainer.make_forward_pass(batch)

            if trainer.grad_acc_steps > 1:
                loss = loss / trainer.grad_acc_steps
            loss.backward()
            if (step + 1) % trainer.grad_acc_steps == 0:
                torch.nn.utils.clip_grad_norm_(
                    trainer.model.parameters(), trainer.params["max_grad_norm"]
                )
                noise_function = trainer.params["noise_approach"]
                if noise_function == "anticorrelated_noise_prev_term":
                    noise.add_anticorrelated_noise_prev_term(optimizer, device)
                if noise_function == "gausian_noise":
                    noise.add_gausian_noise(optimizer, device)
                if noise_function == "anticorrelated_noise_gradient":
                    noise.add_anticorrelated_noise_gradient(optimizer, device)
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()

            if (step+1)% 200 == 0:
                results, mrr = trainer.evaluate(test_data,30,10)
                print("recall at 10:"+str(results))
                print("mrr: "+str(mrr))
        print("Start evaluation after epoch: " + str(e))
        trainer.model.eval()
        results, mrr = trainer.evaluate(test_data,30,10)
        # index, mrr = evaluator.evaluate_mrr(trainer.model)
        print("---------------------------Results in Epoch------------------------:" + str(e))
        print(results)
        # Recall writing in a file
        f = open(trainer.params["training_result_update_file"], 'a+')
        f.write("Results in Epoch " + str(e) + ' : ' + str(results) + '\n')
        f.close()
        # writing mrrs
        f1 = open('Results_Mrr.txt', 'a+')
        f1.write("Results in Epoch " + str(e) + ' : ' + str(mrr) + '\n')
        f1.close()
        #encoding_map = encode_documents(documents, trainer.model, trainer.collator)
        epoch_output_folder_path = os.path.join(
            trainer.params["model_dump_folder"], "epoch_{}".format(e)
        )
        save_model(trainer.model,trainer.tokenizer,  epoch_output_folder_path)
        trainer.model.train()


train(10)