Class DualEncoderDataLoader

data_loaders/DualEncoderDataLoader.py:11–165

import random

class DualEncoderDataLoader(pl.LightningDataModule):
  def __init__(self, tokenizer, data_file, batch_size):
    super().__init__()
    self.tokenizer = tokenizer
    self.data_file = data_file
    self.batch_size = batch_size

  def encode_sentences(self, positives, negatives, max_length=512, pad_to_max_length=False, return_tensors="pt"):
    encoded_dict = self.tokenizer(
            positives,
            max_length=max_length,
            padding="max_length" if pad_to_max_length else "longest",
            truncation=True,
            return_tensors=return_tensors
        )

    pos_input_ids = encoded_dict['input_ids']
    pos_attention_masks = encoded_dict['attention_mask']

    negatives_unstacked = list(np.reshape(negatives, -1))

    encoded_dict = self.tokenizer(
            negatives_unstacked,
            max_length=max_length,
            padding="max_length" if pad_to_max_length else "longest",
            truncation=True,
            return_tensors=return_tensors
        )

    neg_input_ids = encoded_dict['input_ids'].view(len(negatives), len(negatives[0]), -1)
    neg_attention_masks = encoded_dict['attention_mask'].view(len(negatives), len(negatives[0]), -1)
  
    batch = {
        "pos_input_ids": pos_input_ids,
        "pos_attention_masks": pos_attention_masks,
        "neg_input_ids": neg_input_ids,
        "neg_attention_masks": neg_attention_masks,
    }

    return batch

  def triple_to_string(self, x):
    return " </s> ".join([item.strip() for item in x])
  
  def load_tsv_files(self, filepaths):
    srcs = []

    for filepath in filepaths:
        with open(filepath) as tsv:
            for line in tsv:
                parts = line.strip().split("\t")

                if len(parts) == 2:
                    claim = parts[0].split("||")[0].replace("[CLAIM]", "").strip()
                    question = parts[0].split("||")[1].replace("[QUESTION]", "").strip()
                    answer = parts[1].split("||")[1]

                    srcs.append([claim, question, answer])