dev_dpr.py

import os
import json
import random
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from pprint import pprint
from time import strftime
# from dataset import Dataset
from datasets import ( Dataset, DatasetDict,  Features, Value)
import torch
from torch.utils.data import DataLoader, TensorDataset
import torch.nn.functional as F
import pickle

from transformers import (
    AutoTokenizer,
    BertModel, BertPreTrainedModel,
    AdamW, get_linear_schedule_with_warmup,
    TrainingArguments,
)

def set_seed(random_seed):
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed(random_seed)
    torch.cuda.manual_seed_all(random_seed)  # if use multi-GPU
    random.seed(random_seed)
    np.random.seed(random_seed)

# Anwer
class DenseRetrieval:

    def __init__(self, args, dataset, num_neg, tokenizer, p_encoder, q_encoder):

        '''
        학습과 추론에 사용될 여러 셋업을 마쳐봅시다.
        '''

        self.args = args
        self.dataset = dataset
        self.num_neg = num_neg

        self.tokenizer = tokenizer
        self.p_encoder = p_encoder
        self.q_encoder = q_encoder

        self.prepare_in_batch_negative(num_neg=num_neg)

    def prepare_in_batch_negative(self, dataset=None, num_neg=2, tokenizer=None):

        if dataset is None:
            dataset = self.dataset

        if tokenizer is None:
            tokenizer = self.tokenizer

        # 1. In-Batch-Negative 만들기
        # CORPUS를 np.array로 변환해줍니다.
        corpus = np.array(list(set([example for example in dataset['context']])))
        p_with_neg = []

        for c in tqdm(dataset['context'],  desc='In-batch-Negative'):

            while True:
                neg_idxs = np.random.randint(len(corpus), size=num_neg)

                if not c in corpus[neg_idxs]:
                    p_neg = corpus[neg_idxs]

                    p_with_neg.append(c)
                    p_with_neg.extend(p_neg)
                    break

        # 2. (Question, Passage) 데이터셋 만들어주기
        q_seqs = tokenizer(dataset['question'].tolist(), padding="max_length", truncation=True, return_tensors='pt')
        p_seqs = tokenizer(p_with_neg, padding="max_length", truncation=True, return_tensors='pt')

        max_len = p_seqs['input_ids'].size(-1)
        p_seqs['input_ids'] = p_seqs['input_ids'].view(-1, num_neg+1, max_len)
        p_seqs['attention_mask'] = p_seqs['attention_mask'].view(-1, num_neg+1, max_len)
        p_seqs['token_type_ids'] = p_seqs['token_type_ids'].view(-1, num_neg+1, max_len)

        train_dataset = TensorDataset(
            p_seqs['input_ids'], p_seqs['attention_mask'], p_seqs['token_type_ids'],
            q_seqs['input_ids'], q_seqs['attention_mask'], q_seqs['token_type_ids']
        )

        self.train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=self.args.per_device_train_batch_size)

        


    def train(self, args=None):

        if args is None:
            args = self.args
        batch_size = args.per_device_train_batch_size

        # Optimizer
        no_decay = ['bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in self.p_encoder.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
            {'params': [p for n, p in self.p_encoder.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},
            {'params': [p for n, p in self.q_encoder.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': args.weight_decay},
            {'params': [p for n, p in self.q_encoder.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
        ]
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
        t_total = len(self.train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)

        
        # Start training!
        global_step = 0

        self.p_encoder.zero_grad()
        self.q_encoder.zero_grad()
        torch.cuda.empty_cache()

        tm = strftime('%Y.%m.%d - %H:%M:%S')
        save_dir = self.args.output_dir + f'/{tm}'
        os.makedirs(save_dir)

        train_iterator = tqdm(range(int(args.num_train_epochs)), desc="Epoch")
        # for _ in range(int(args.num_train_epochs)):
        for _ in train_iterator:

            with tqdm(self.train_dataloader, unit="batch") as tepoch:
                for batch in tepoch:

                    p_encoder.train()
                    q_encoder.train()
                    batch_size = batch[0].shape[0]
                    targets = torch.zeros(batch_size).long() # positive example은 전부 첫 번째에 위치하므로
                    targets = targets.to(args.device)

                    p_inputs = {
                        'input_ids': batch[0].view(batch_size * (self.num_neg + 1), -1).to(args.device),
                        'attention_mask': batch[1].view(batch_size * (self.num_neg + 1), -1).to(args.device),
                        'token_type_ids': batch[2].view(batch_size * (self.num_neg + 1), -1).to(args.device)
                    }

                    q_inputs = {
                        'input_ids': batch[3].to(args.device),
                        'attention_mask': batch[4].to(args.device),
                        'token_type_ids': batch[5].to(args.device)
                    }

                    p_outputs = self.p_encoder(**p_inputs)  # (batch_size*(num_neg+1), emb_dim)
                    q_outputs = self.q_encoder(**q_inputs)  # (batch_size*, emb_dim)

                    # Calculate similarity score & loss
                    
                    p_outputs = p_outputs.view(batch_size, self.num_neg + 1, -1)
                    q_outputs = q_outputs.view(batch_size, 1, -1)
                    

                    sim_scores = torch.bmm(q_outputs, torch.transpose(p_outputs, 1, 2)).squeeze()  #(batch_size, num_neg + 1)
                    sim_scores = sim_scores.view(batch_size, -1)
                    sim_scores = F.log_softmax(sim_scores, dim=1)

                    loss = F.nll_loss(sim_scores, targets)
                    tepoch.set_postfix(loss=f'{str(loss.item())}')

                    loss.backward() 
                    optimizer.step()
                    scheduler.step()

                    self.p_encoder.zero_grad()
                    self.q_encoder.zero_grad()

                    global_step += 1

                    torch.cuda.empty_cache()

                    del p_inputs, q_inputs
            save_path = save_dir + f'/epochs:{_}'
            os.makedirs(save_path)
            
            torch.save(p_encoder, save_path + f'/p_encoder.pt')
            torch.save(q_encoder, save_path + f'/q_encoder.pt')

    # def get_dense_embedding(self) -> None:

    #     """
    #     Summary:
    #         Passage Embedding을 만들고
    #         TFIDF와 Embedding을 pickle로 저장합니다.
    #         만약 미리 저장된 파일이 있으면 저장된 pickle을 불러옵니다.
    #     """

    #     # Pickle을 저장합니다.
    #     pickle_name = f"dense_embedding.bin"
    #     p_dense_name = f"p_dense.bin"
    #     emd_path = os.path.join(self.data_path, pickle_name)
    #     dense_path = os.path.join(self.data_path, p_dense_name)

    #     if os.path.isfile(emd_path) and os.path.isfile(dense_path):
    #         with open(emd_path, "rb") as file:
    #             self.p_embedding = pickle.load(file)
    #         with open(dense_path, "rb") as file:
    #             self.tfidfv = pickle.load(file)
    #         print("Embedding pickle load.")
    #     else:
    #         print("Build passage embedding")
    #         self.p_embedding = self.tfidfv.fit_transform(self.contexts)
    #         print(self.p_embedding.shape)
    #         with open(emd_path, "wb") as file:
    #             pickle.dump(self.p_embedding, file)
    #         with open(dense_path, "wb") as file:
    #             pickle.dump(self.tfidfv, file)
    #         print("Embedding pickle saved.")

    def get_relevant_doc(self, query, k=1, args=None, p_encoder=None, q_encoder=None):
        # if dataset is None:
        dataset = pd.read_json('/opt/ml/input/code/level2_nlp_mrc-nlp-08/csv_data/wikipedia_documents.json')
        dataset = dataset.transpose()
        dataset = dataset[:-1]
        if args is None:
            args = self.args

        if p_encoder is None:
            p_encoder = self.p_encoder

        if q_encoder is None:
            q_encoder = self.q_encoder


        valid_seqs = self.tokenizer(dataset['text'].tolist(), padding="max_length", truncation=True, return_tensors='pt')
        ##
        # valid_seqs = {'input_ids':[], 'attention_mask':[],'token_type_ids':[] }

        # data_context = dataset['context'].tolist()

        # for idx, context in enumerate(tqdm(data_context, desc = 'tokenizing')):
        #     token = tokenizer(context, padding="max_length", truncation=True, return_tensors='pt')
        #     for key, value in token.items():
        #         valid_seqs[key].append(value)
        
        # for key in valid_seqs.keys():
        #     valid_seqs[key] =  torch.cat(valid_seqs[key])

        passage_dataset = TensorDataset(
            valid_seqs['input_ids'], valid_seqs['attention_mask'], valid_seqs['token_type_ids']
        )
        self.passage_dataloader = DataLoader(passage_dataset, batch_size=self.args.per_device_train_batch_size)

        
        batch_size = args.per_device_train_batch_size
        with torch.no_grad():
            p_encoder.eval()
            q_encoder.eval()

            # passage
            p_embs = []
            for batch in tqdm(self.passage_dataloader):

                batch = tuple(t.to(args.device) for t in batch)
                p_inputs = {
                    'input_ids': batch[0],
                    'attention_mask': batch[1],
                    'token_type_ids': batch[2]
                }
                p_emb = p_encoder(**p_inputs).to('cpu')
                p_embs.append(p_emb)

        p_embs = torch.stack(p_embs, dim=0).view(len(self.passage_dataloader.dataset), -1)  # (num_passage, emb_dim)
        
        #question
        # q_seqs_val = self.tokenizer([query['question']], padding="max_length", truncation=True, return_tensors='pt').to(args.device)
        # q_emb = q_encoder(**q_seqs_val).to('cpu')  # (num_query=1, emb_dim)

        answer = []
        for i in tqdm(range(len(query)), desc='q'):
            q_token = self.tokenizer([query['question'][i]], padding="max_length", truncation=True, return_tensors='pt').to(args.device)
            q_emb = q_encoder(**q_token).to('cpu')  # (num_query=1, emb_dim)

            dot_prod_scores = torch.matmul(q_emb, torch.transpose(p_embs, 0, 1))
            rank = torch.argsort(dot_prod_scores, dim=1, descending=True).squeeze()
            topk_context = rank[:k].tolist()
            join_context = ' '.join([ dataset['text'][i] for i in topk_context])
            tmp =  {
                    # Query와 해당 id를 반환합니다.
                    "question": query["question"][i],
                    "id": query["id"][i],
                    # Retrieve한 Passage의 id, context를 반환합니다.
                    "context": join_context
                }
            
            answer.append(tmp)
        
        dqas = pd.DataFrame(answer)
        f = Features(
            {
                "context": Value(dtype="string", id=None),
                "id": Value(dtype="string", id=None),
                "question": Value(dtype="string", id=None),
            }
        )

        datasets = DatasetDict({"validation": Dataset.from_pandas(dqas, features=f)})
        return  datasets
    

class BertEncoder(BertPreTrainedModel):

    def __init__(self, config):
        super(BertEncoder, self).__init__(config)

        self.bert = BertModel(config)
        self.init_weights()


    def forward(
            self,
            input_ids,
            attention_mask=None,
            token_type_ids=None
        ):

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids
        )

        pooled_output = outputs[1]
        return pooled_output
    

if __name__ == "__main__":
    set_seed(42)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    # 데이터셋과 모델은 아래와 같이 불러옵니다.
    data_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), "csv_data")
    train_dataset = pd.read_csv(os.path.join(data_path, "squad_kor_v1.csv"))
    train_dataset = train_dataset
    output_dir=os.path.join(os.path.abspath(os.path.dirname(__file__)), "dense_retireval")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    

    args = TrainingArguments(
        output_dir=output_dir,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        learning_rate=3e-4,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        num_train_epochs=3,
        weight_decay=0.01
    )
    model_checkpoint = 'klue/bert-base'
    # model_checkpoint = "klue/roberta-large"

    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    p_encoder = BertEncoder.from_pretrained(model_checkpoint).to(args.device)
    q_encoder = BertEncoder.from_pretrained(model_checkpoint).to(args.device)
    retriever = DenseRetrieval(args=args, dataset=train_dataset, num_neg=3, tokenizer=tokenizer, p_encoder=p_encoder, q_encoder=q_encoder)
    retriever.train()

    # query = '제주도 시청의 주소는 뭐야?'
    # results = retriever.get_relevant_doc(query=query, k=5)
    # print(f"[Search Query] {query}\n")

    # indices = results.tolist()
    # for i, idx in enumerate(indices):
    #     print(f"Top-{i + 1}th Passage (Index {idx})")
    #     pprint(retriever.dataset['context'][idx])