cal_persona_label.py

import json
import pickle
from sys import argv

import torch
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
from tqdm import tqdm
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers.data.processors.utils import InputExample

NLI_MODEL_PATH = './persona_nli'

# The original train file
input_file = argv[1]
# The output file that saves the NLI logits given the train samples
output_file = argv[2]
bz = int(argv[3].strip())
gpu = argv[4].strip()

def get_dataloader(input_examples, tokenizer, device, batch_size=512):
    features = convert_examples_to_features(
        input_examples,
        tokenizer,
        label_list=['0', '1'],
        max_length=128,
        output_mode='classification',
    )
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long).to(device)
    all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long).to(device)
    dataset = TensorDataset(all_input_ids, all_attention_mask)
    dataloader = DataLoader(dataset, batch_size=batch_size)
    return dataloader


def read_txt_data(input_file):
    with open(input_file, 'r') as f:
        lines = f.readlines()
    # for two agent
    cur_persona1, cur_persona2 = [], []
    cur_dialogs1, cur_dialogs2 = [], []
    personas, dialogs = [], []

    start = True
    for line in lines:
        if 'your persona:' in line or 'partner\'s persona' in line:
            if start and len(cur_persona1) > 0: # start a new dialogue
                personas.append(cur_persona1)
                personas.append(cur_persona2)
                dialogs.append(cur_dialogs1)
                dialogs.append(cur_dialogs2)
                cur_persona1, cur_persona2 = [], []
                cur_dialogs1, cur_dialogs2 = [], []
            start = False
            # parsing persona
            if 'your persona:' in line: 
                persona_index = line.find('your persona:')
                persona = line[persona_index + 14: -1]
                cur_persona1.append(persona)
            elif 'partner\'s persona' in line:
                persona_index = line.find('partner\'s persona:')
                persona = line[persona_index + 19: -1]
                cur_persona2.append(persona)
        else:
            start = True
            space_index = line.find(' ')
            sents = line[space_index + 1:].split('\t')
            cur_dialogs1.append(sents[1])
            cur_dialogs2.append(sents[0])

    return personas, dialogs


def read_json_data(input_file):
    with open(input_file, 'r') as f:
        data = json.load(f)
    examples = []
    cnt = 0
    for d in data:
        examples.append(InputExample(str(cnt), d[0], d[2], '0'))
        cnt += 1
    return examples

# load tokenizer and model (single gpu is enough)
tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL_PATH)
device = torch.device(f'cuda:{gpu}') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
model.eval()

# make prediction on persona data
pred_results = []
if '.json' in input_file: # json file 
    all_logits = None
    input_examples = read_json_data(input_file)
    train_dataloader = get_dataloader(input_examples, tokenizer, device, bz)
    with torch.no_grad():
        for batch in tqdm(train_dataloader):
            inputs = {"input_ids": batch[0], "attention_mask": batch[1]}
            outputs = model(**inputs)
            if all_logits is None:
                all_logits = outputs[0].cpu().detach()
            else:
                all_logits = torch.cat((all_logits, outputs[0].cpu().detach()), dim=0)
    print (all_logits.shape) # [n, 3] ???
    all_logits = all_logits.numpy()
    with open(output_file, 'wb') as f:
        pickle.dump(all_logits, f)
else: # txt file (persona raw text file)
    personas, dialogs = read_txt_data(input_file)
    entailed_results = []
    with torch.no_grad():
        for i in tqdm(range(len(personas))): #  for every dialog
            cur_persona = personas[i]
            cur_dialogs = dialogs[i]
            cnt = 0
            cur_pred_results = []
            for persona in cur_persona: # for every persona entry
                input_examples = []
                for dialog in cur_dialogs: # for every uttr
                    input_examples.append(InputExample(str(cnt), persona, dialog, '0')) # A single training/test example for simple sequence classification
                    cnt += 1
                train_dataloader = get_dataloader(input_examples, tokenizer, device)
                all_logits = None
                for batch in train_dataloader:
                    inputs = {"input_ids": batch[0], "attention_mask": batch[1]}
                    outputs = model(**inputs)
                    if all_logits is None: 
                        all_logits = outputs[0].detach()
                    else:
                        all_logits = torch.cat((all_logits, outputs[0].detach()), dim=0)
                # dimension: [n, 3] -> [n]
                results = torch.argmax(all_logits, dim=1)
                for id, res in enumerate(results): # a persona entry may entails None, one or severl dialog responses
                    if res == 2: # entail
                        entailed_results.append((persona, cur_dialogs[id]))
                cur_pred_results.append(all_logits.cpu())
            # dimension: num_dialogs x matched persona-uttr pair in the each dialog
            pred_results.append(cur_pred_results) 
        
    
    torch.save(pred_results, f'{output_file}/entailment_scores.bin')
    with open(f'{output_file}/entailed_sentences.json', 'w') as f:
        json.dump(entailed_results, f)