|
|
import torch |
|
|
import random |
|
|
from typing import Any, Dict, Union |
|
|
|
|
|
|
|
|
def create_negative_dataset(train_dataloader): |
|
|
|
|
|
list = [] |
|
|
for step, inputs in enumerate(train_dataloader): |
|
|
input_ids, input_masks, segment_ids, label_ids = inputs |
|
|
input_ids = input_ids.tolist() |
|
|
input_masks = input_masks.tolist() |
|
|
segment_ids = segment_ids.tolist() |
|
|
label_ids = label_ids.tolist() |
|
|
|
|
|
for i in range(len(input_ids)-1): |
|
|
input_id = input_ids[i] |
|
|
input_mask = input_masks[i] |
|
|
segment_id = segment_ids[i] |
|
|
label_id = label_ids[i] |
|
|
batch_dict = {"labels": label_id, "input_ids": input_id, "token_type_ids": segment_id, |
|
|
"attention_mask": input_mask} |
|
|
list.append(batch_dict) |
|
|
|
|
|
negative_dataset = {} |
|
|
|
|
|
for line in list: |
|
|
label = int(line["labels"]) |
|
|
inputs = line |
|
|
inputs.pop("labels") |
|
|
if label not in negative_dataset.keys(): |
|
|
negative_dataset[label] = [inputs] |
|
|
else: |
|
|
negative_dataset[label].append(inputs) |
|
|
|
|
|
return negative_dataset |
|
|
|
|
|
|
|
|
def generate_positive_sample(negative_data, args, label: torch.Tensor): |
|
|
positive_num = args.positive_num |
|
|
|
|
|
positive_sample = [] |
|
|
for index in range(label.shape[0]): |
|
|
input_label = int(label[index]) |
|
|
positive_sample.extend(random.sample(negative_data[input_label], positive_num)) |
|
|
|
|
|
return list_item_to_tensor(positive_sample) |
|
|
|
|
|
def _prepare_inputs(device, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]: |
|
|
for k, v in inputs.items(): |
|
|
if isinstance(v, torch.Tensor): |
|
|
inputs[k] = v.to(device) |
|
|
|
|
|
return inputs |
|
|
|
|
|
|
|
|
def list_item_to_tensor(inputs_list): |
|
|
batch_list = {} |
|
|
for key, value in inputs_list[0].items(): |
|
|
batch_list[key] = [] |
|
|
for inputs in inputs_list: |
|
|
for key, value in inputs.items(): |
|
|
batch_list[key].append(value) |
|
|
|
|
|
batch_tensor = {} |
|
|
for key, value in batch_list.items(): |
|
|
batch_tensor[key] = torch.tensor(value) |
|
|
return batch_tensor |
|
|
|
