<filename>src/dataloader.py
import torch
import torch.nn as nn
import torch.utils.data as data
from torch.utils.data import DataLoader

import math
from tqdm import tqdm
import random
import logging

logger = logging.getLogger()

# from transformers import BertTokenizer
# bert_tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
from src.config import get_params

params = get_params()
from transformers import AutoTokenizer

auto_tokenizer = AutoTokenizer.from_pretrained(params.model_name)
pad_token_label_id = nn.CrossEntropyLoss().ignore_index


# domain2labels = {"politics": politics_labels, "science": science_labels, "music": music_labels, "literature": literature_labels, "ai": ai_labels }
# "drugs":conll_to_drugs, "single_drugs":single_drugs, "conll_to_drugs":conll_to_drugs, "wnut_to_drugs":wnut_to_drugs, "btc_to_drugs":btc_to_drugs

def remove_duplicates_with_order(seq):
    seen = set()
    seen_add = seen.add
    return [x for x in seq if not (x in seen or seen_add(x))]

def get_domain2labels(tgt_dm, src_dm):
    politics_labels = ['O', 'B-country', 'B-politician', 'I-politician', 'B-election', 'I-election', 'B-person',
                       'I-person', 'B-organisation', 'I-organisation', 'B-location', 'B-misc', 'I-location',
                       'I-country', 'I-misc', 'B-politicalparty', 'I-politicalparty', 'B-event', 'I-event']
    science_labels = ['O', 'B-scientist', 'I-scientist', 'B-person', 'I-person', 'B-university', 'I-university',
                      'B-organisation', 'I-organisation', 'B-country', 'I-country', 'B-location', 'I-location',
                      'B-discipline', 'I-discipline', 'B-enzyme', 'I-enzyme', 'B-protein', 'I-protein',
                      'B-chemicalelement', 'I-chemicalelement', 'B-chemicalcompound', 'I-chemicalcompound',
                      'B-astronomicalobject', 'I-astronomicalobject', 'B-academicjournal', 'I-academicjournal',
                      'B-event', 'I-event', 'B-theory', 'I-theory', 'B-award', 'I-award', 'B-misc', 'I-misc']
    music_labels = ['O', 'B-musicgenre', 'I-musicgenre', 'B-song', 'I-song', 'B-band', 'I-band', 'B-album', 'I-album',
                    'B-musicalartist', 'I-musicalartist', 'B-musicalinstrument', 'I-musicalinstrument', 'B-award',
                    'I-award', 'B-event', 'I-event', 'B-country', 'I-country', 'B-location', 'I-location',
                    'B-organisation', 'I-organisation', 'B-person', 'I-person', 'B-misc', 'I-misc']
    literature_labels = ["O", "B-book", "I-book", "B-writer", "I-writer", "B-award", "I-award", "B-poem", "I-poem",
                         "B-event", "I-event", "B-magazine", "I-magazine", "B-literarygenre", "I-literarygenre",
                         'B-country', 'I-country', "B-person", "I-person", "B-location", "I-location", 'B-organisation',
                         'I-organisation', 'B-misc', 'I-misc']
    ai_labels = ["O", "B-field", "I-field", "B-task", "I-task", "B-product", "I-product", "B-algorithm", "I-algorithm",
                 "B-researcher", "I-researcher", "B-metrics", "I-metrics", "B-programlang", "I-programlang",
                 "B-conference", "I-conference", "B-university", "I-university", "B-country", "I-country", "B-person",
                 "I-person", "B-organisation", "I-organisation", "B-location", "I-location", "B-misc", "I-misc"]

    single_wnut = ["O", "B-creative-work", "I-creative-work", "B-person", "I-person", "B-corporation", "I-corporation",
                   "B-location", "I-location", "B-group", "I-group", "B-product", "I-product"]  # 12 without "O"
    single_conll = ['O', 'B-organisation', 'I-organisation', 'B-misc', 'I-misc', 'B-person', 'I-person', 'B-location',
                    'I-location']  # 8 without "O"
    single_btc = ["O", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG", "I-ORG"]  # 6 without "O"
    single_drugs = ["O", "B-Drug", "I-Drug"]  # 2 without "O"

    conll_to_drugs = ['O', 'B-organisation', 'I-organisation', 'B-misc', 'I-misc', 'B-person', 'I-person', 'B-location',
                      'I-location', 'B-Drug', 'I-Drug']
    wnut_to_drugs = ["O", "B-creative-work", "I-creative-work", "B-person", "I-person", "B-corporation",
                     "I-corporation", "B-location", "I-location", "B-product", "I-product", "B-Drug", "I-Drug"]
    btc_to_drugs = ["O", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG", "I-ORG", "B-Drug", "I-Drug"]
    wnut_to_btc = ["O", "B-creative-work", "I-creative-work", "B-person", "I-person", "B-corporation", "I-corporation",
                   "B-location", "I-location", "B-product", "I-product", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG",
                   "I-ORG", "B-Drug", "I-Drug"]

    domain2labels = {"politics": politics_labels, "science": science_labels, "music": music_labels,
                     "literature": literature_labels, "ai": ai_labels}
    custom_domains = {'wnut': single_wnut, 'nutot': single_wnut, 'btc': single_btc, 'drugs': single_drugs,
                      'wiki': single_drugs, 'conll': single_conll}

    # create custom entity class set
    if tgt_dm not in domain2labels:
        if tgt_dm in custom_domains and src_dm in custom_domains:
            domain2labels[tgt_dm] = custom_domains[tgt_dm]
            domain2labels[tgt_dm].extend(custom_domains[src_dm])
            domain2labels[tgt_dm] = remove_duplicates_with_order(domain2labels[tgt_dm])  # remove duplicate entity classes
        else:
            raise Exception("Target or Source Domain not in Custom Domain Set! Fix that")
    elif params.tgt_dm not in domain2labels and params.src_dm in custom_domains and src_dm != 'conll':
        raise Exception("If you want to use a standard domain with custom source domain --> you gotta configure this.")

    return domain2labels

def split_text(token_list, label_list, cut_strat):
    dot_enc = auto_tokenizer.convert_tokens_to_ids('.')
    splitted_token_lists, splitted_label_lists = [], []
    amt_cutoff=0
    amt_total_splits=1


    if cut_strat != 1: # if cut_strat is 1 than directly cutoff long texts.
        if cut_strat == 2:
            start_at_indx = 400
            cut_until = 500
        elif cut_strat == 3:
            start_at_indx = 0#Cut into sentences
            cut_until = 0
        elif cut_strat == 4: #Ignore longer texts
            return splitted_token_lists, splitted_label_lists, len(token_list), amt_total_splits
        else:
            print('Cut Strategy: ' , str(cut_strat))
            raise ValueError
        while len(token_list) > cut_until and start_at_indx >= 0:
            try:
                cut_at = token_list.index(dot_enc, start_at_indx, 500)
                splitted_token_lists.append(token_list[:cut_at])
                splitted_label_lists.append(label_list[:cut_at])
                if cut_at < len(token_list):
                    token_list = token_list[cut_at+1:]
                    label_list = label_list[cut_at+1:]
                    amt_total_splits+=1
                else:
                    token_list = []
                    label_list = []

                if cut_strat == 2:#Reset starting point for next possible split
                    start_at_indx = 400
            except ValueError:
                logger.info("Split not found enlarging search area")
                start_at_indx = start_at_indx - 100 #If no dot is in that area --> make search are bigger or continue

    if len(token_list) > 500:
        amt_cutoff = len(token_list) - 500
        splitted_token_lists.append(token_list[:500])
        splitted_label_lists.append(label_list[:500])
    else:
        splitted_token_lists.append(token_list)
        splitted_label_lists.append(label_list)

    return splitted_token_lists, splitted_label_lists, amt_cutoff, amt_total_splits




def read_ner(datapath, tgt_dm, src_dm):
    inputs, labels = [], []
    domain2labels = get_domain2labels(tgt_dm, src_dm)
    amt_cut_off=0
    amt_tokens=0
    amt_splits=0

    with open(datapath, "r") as fr:
        token_list, label_list  = [], []

        for i, line in enumerate(fr):
            line = line.strip()
            if line == "":
                if len(token_list) > 0:
                    assert len(token_list) == len(label_list)

                    if len(token_list) > 500 or params.cut_strategy==3:
                        splitted_token_lists, splitted_label_lists, recent_cutoff, rec_splits = split_text(token_list, label_list, params.cut_strategy)
                        amt_cut_off += recent_cutoff
                        amt_splits += rec_splits



                        for cut_tokens in splitted_token_lists:
                            inputs.append([auto_tokenizer.cls_token_id] + cut_tokens + [auto_tokenizer.sep_token_id])
                        for cut_labels in splitted_label_lists:
                            labels.append([pad_token_label_id] + cut_labels + [pad_token_label_id])
                    else:
                        amt_tokens += len(token_list)
                        inputs.append([auto_tokenizer.cls_token_id] + token_list + [auto_tokenizer.sep_token_id])
                        labels.append([pad_token_label_id] + label_list + [pad_token_label_id])

                token_list, label_list = [], []
                continue

            splits = line.split("\t")

            if len(splits) != 2:
                print('i:%d, text:%s'%(i, line))

            token = splits[0]
            #Check if its wnut, since this dataset has multiple labels on the testset..
            if ( src_dm != 'wnut' and src_dm != 'nutot' and tgt_dm != 'wnut'  and tgt_dm != 'nutot' ) or "," not in splits[1]:
                label = splits[1]
            else: # If its WNUT and it has commas in the label --> prioritize product, but acutally doesn't matter much.#TODO check if I should build this in properly
                possible_labels = splits[1].split(',')
                if 'B-product' in possible_labels:
                    label = 'B-product'
                elif 'I-product' in possible_labels:
                    label = 'I-product'
                else:
                    label = possible_labels[0]


            subs_ = auto_tokenizer.tokenize(token)
            if len(subs_) > 0:
                label_list.extend([domain2labels[tgt_dm].index(label)] + [pad_token_label_id] * (len(subs_) - 1))
                token_list.extend(auto_tokenizer.convert_tokens_to_ids(subs_))
            else:
                print("length of subwords for %s is zero; its label is %s" % (token, label))

    logger.info('Read total of %d tokens and cutoff %d tokens for file %s and split/cut at least into %d times'%(amt_tokens, amt_cut_off, datapath, amt_splits))
    return inputs, labels


def read_ner_for_bilstm(datapath, tgt_dm, vocab):
    domain2labels = get_domain2labels(tgt_dm, 'Error')
    inputs, labels = [], []
    with open(datapath, "r") as fr:
        token_list, label_list = [], []
        for i, line in enumerate(fr):
            line = line.strip()
            if line == "":
                if len(token_list) > 0:
                    assert len(token_list) == len(label_list)
                    inputs.append(token_list)
                    labels.append(label_list)

                token_list, label_list = [], []
                continue

            splits = line.split("\t")
            token = splits[0]
            label = splits[1]

            token_list.append(vocab.word2index[token])
            label_list.append(domain2labels[tgt_dm].index(label))

    return inputs, labels


class Dataset(data.Dataset):
    def __init__(self, inputs, labels):
        self.X = inputs
        self.y = labels

    def __getitem__(self, index):
        return self.X[index], self.y[index]

    def __len__(self):
        return len(self.X)


PAD_INDEX = 0


class Vocab():
    def __init__(self):
        self.word2index = {"PAD": PAD_INDEX}
        self.index2word = {PAD_INDEX: "PAD"}
        self.n_words = 1

    def index_words(self, word_list):
        for word in word_list:
            if word not in self.word2index:
                self.word2index[word] = self.n_words
                self.index2word[self.n_words] = word
                self.n_words += 1


def get_vocab(path):
    vocabulary = []
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            vocabulary.append(line)
    return vocabulary


def collate_fn(data):
    X, y = zip(*data)
    lengths = [len(bs_x) for bs_x in X]
    max_lengths = max(lengths)
    padded_seqs = torch.LongTensor(len(X), max_lengths).fill_(auto_tokenizer.pad_token_id)
    padded_y = torch.LongTensor(len(X), max_lengths).fill_(pad_token_label_id)
    for i, (seq, y_) in enumerate(zip(X, y)):
        length = lengths[i]
        padded_seqs[i, :length] = torch.LongTensor(seq)
        padded_y[i, :length] = torch.LongTensor(y_)

    return padded_seqs, padded_y


def collate_fn_for_bilstm(data):
    X, y = zip(*data)
    lengths = [len(bs_x) for bs_x in X]
    max_lengths = max(lengths)
    padded_seqs = torch.LongTensor(len(X), max_lengths).fill_(PAD_INDEX)
    for i, seq in enumerate(X):
        length = lengths[i]
        padded_seqs[i, :length] = torch.LongTensor(seq)

    lengths = torch.LongTensor(lengths)
    return padded_seqs, lengths, y


def get_dataloader_for_bilstmtagger(params):
    vocab_src = get_vocab("ner_data/conll2003/vocab.txt")
    vocab_tgt = get_vocab("ner_data/%s/vocab.txt" % params.tgt_dm)
    vocab = Vocab()
    vocab.index_words(vocab_src)
    vocab.index_words(vocab_tgt)

    logger.info("Load training set data ...")
    conll_inputs_train, conll_labels_train = read_ner_for_bilstm("ner_data/conll2003/train.txt", params.tgt_dm, vocab)
    inputs_train, labels_train = read_ner_for_bilstm("ner_data/%s/train.txt" % params.tgt_dm, params.tgt_dm, vocab)
    inputs_train = inputs_train * 10 + conll_inputs_train
    labels_train = labels_train * 10 + conll_labels_train

    logger.info("Load dev set data ...")
    inputs_dev, labels_dev = read_ner_for_bilstm("ner_data/%s/dev.txt" % params.tgt_dm, params.tgt_dm, vocab)

    logger.info("Load test set data ...")
    inputs_test, labels_test = read_ner_for_bilstm("ner_data/%s/test.txt" % params.tgt_dm, params.tgt_dm, vocab)

    logger.info("train size: %d; dev size %d; test size: %d;" % (len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=params.batch_size, shuffle=True,
                                  collate_fn=collate_fn_for_bilstm)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=params.batch_size, shuffle=False,
                                collate_fn=collate_fn_for_bilstm)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=params.batch_size, shuffle=False,
                                 collate_fn=collate_fn_for_bilstm)

    return dataloader_train, dataloader_dev, dataloader_test, vocab


def load_corpus(tgt_dm):
    print("Loading corpus ...")
    data_path = "enwiki_corpus/%s_removebracket.tok" % tgt_dm
    sent_list = []
    with open(data_path, "r") as fr:
        for i, line in tqdm(enumerate(fr)):
            line = line.strip()
            sent_list.append(line)
    return sent_list


def get_dataloader(params):
    logger.info("Load training set data")
    inputs_train, labels_train = read_ner("ner_data/%s/train.txt" % params.tgt_dm, params.tgt_dm, params.src_dm)
    domain2labels = get_domain2labels(params.tgt_dm, params.src_dm)

    if params.n_samples != -1:
        logger.info("Few-shot on %d samples" % params.n_samples)
        inputs_train = inputs_train[:params.n_samples]
        labels_train = labels_train[:params.n_samples]
    logger.info("Load development set data")
    inputs_dev, labels_dev = read_ner("ner_data/%s/dev.txt" % params.tgt_dm, params.tgt_dm, params.src_dm)
    logger.info("Load test set data")
    inputs_test, labels_test = read_ner("ner_data/%s/test.txt" % params.tgt_dm, params.tgt_dm, params.src_dm)

    logger.info("label distribution for training set")
    label_distri_train = {}
    count_tok_train = 0
    for label_seq in labels_train:
        for label in label_seq:
            if label != pad_token_label_id:
                label_name = domain2labels[params.tgt_dm][label]
                if "B" in label_name:
                    count_tok_train += 1
                    label_name = label_name.split("-")[1]
                    if label_name not in label_distri_train:
                        label_distri_train[label_name] = 1
                    else:
                        label_distri_train[label_name] += 1
    print(label_distri_train)
    for key in label_distri_train:
        label_distri_train[key] /= count_tok_train
    logger.info(label_distri_train)

    logger.info("label distribution for dev set")
    label_distri_dev = {}
    count_tok_test = 0
    for label_seq in labels_dev:
        for label in label_seq:
            if label != pad_token_label_id:
                label_name = domain2labels[params.tgt_dm][label]
                if "B" in label_name:
                    count_tok_test += 1
                    label_name = label_name.split("-")[1]
                    if label_name not in label_distri_dev:
                        label_distri_dev[label_name] = 1
                    else:
                        label_distri_dev[label_name] += 1
    print(label_distri_dev)
    for key in label_distri_dev:
        label_distri_dev[key] /= count_tok_test
    logger.info(label_distri_dev)

    logger.info("label distribution for test set")
    label_distri_test = {}
    count_tok_test = 0
    for label_seq in labels_test:
        for label in label_seq:
            if label != pad_token_label_id:
                label_name = domain2labels[params.tgt_dm][label]
                if "B" in label_name:
                    count_tok_test += 1
                    label_name = label_name.split("-")[1]
                    if label_name not in label_distri_test:
                        label_distri_test[label_name] = 1
                    else:
                        label_distri_test[label_name] += 1
    print(label_distri_test)
    for key in label_distri_test:
        label_distri_test[key] /= count_tok_test
    logger.info(label_distri_test)

    if params.conll and params.joint: #No worries, this is only done when training is executed in a joint manner. Furthermore our dataset are nearly equal so balancing is not so important.
        conll_inputs_train, conll_labels_train = read_ner("ner_data/conll2003/train.txt", params.tgt_dm, params.src_dm)
        inputs_train = inputs_train * 50  # augment the target domain data to balance the source and target domain data
        labels_train = labels_train * 50
        inputs_train = inputs_train + conll_inputs_train
        labels_train = labels_train + conll_labels_train

    logger.info("train size: %d; dev size %d; test size: %d;" % (len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=params.batch_size, shuffle=True,
                                  collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=params.batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=params.batch_size, shuffle=False,
                                 collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_conll2003_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/conll2003/train.txt", params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/conll2003/dev.txt", params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/conll2003/test.txt", params.tgt_dm, params.src_dm)

    logger.info("conll2003 dataset: train size: %d; dev size %d; test size: %d" % (
    len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_btc_dataloader(batch_size, params):
    # According to https://github.com/GateNLP/broad_twitter_corpus
    # Training: everything else
    inputs_train1, labels_train1 = read_ner("ner_data/btc/a.conll", params.tgt_dm, params.src_dm)
    inputs_train2, labels_train2 = read_ner("ner_data/btc/b.conll", params.tgt_dm, params.src_dm)
    inputs_train3, labels_train3 = read_ner("ner_data/btc/e.conll", params.tgt_dm, params.src_dm)
    inputs_train4, labels_train4 = read_ner("ner_data/btc/g.conll", params.tgt_dm, params.src_dm)
    inputs_train5, labels_train5 = read_ner("ner_data/btc/h.conll", params.tgt_dm, params.src_dm)

    # Development: second half of Section H
    inputs_dev = inputs_train5[math.floor( len(inputs_train5) / 2):]
    labels_dev = labels_train5[math.floor( len(labels_train5) / 2):]
    inputs_train5 = inputs_train5[:math.floor( len(inputs_train5) / 2)]
    labels_train5 = labels_train5[:math.floor( len(labels_train5) / 2)]

    # Test: Section F
    inputs_test, labels_test = read_ner("ner_data/btc/f.conll", params.tgt_dm, params.src_dm)

    # Aggregate Train sets
    inputs_train = inputs_train1 + inputs_train2 + inputs_train3 + inputs_train4 + inputs_train5
    labels_train = labels_train1 + labels_train2 + labels_train3 + labels_train4 + labels_train5

    logger.info("Broad Twitter Corpus dataset: train size: %d; dev size %d; test size: %d" % (
        len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_wnut_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/wnut/train.txt", params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/wnut/dev.txt", params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/wnut/test.txt", params.tgt_dm, params.src_dm)

    logger.info("W-NUT 2017 dataset: train size: %d; dev size %d; test size: %d" % (
        len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_generic_source_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/%s/train.txt" % params.src_dm, params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/%s/dev.txt" % params.src_dm, params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/%s/test.txt" % params.src_dm, params.tgt_dm, params.src_dm)

    logger.info("%s dataset: train size: %d; dev size %d; test size: %d" % (params.src_dm,
                                                                            len(inputs_train), len(inputs_dev),
                                                                            len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_wiki_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/wiki/train.txt", params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/wiki/dev.txt", params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/wiki/test.txt", params.tgt_dm, params.src_dm)

    logger.info("Wikipedia distantly supervised dataset: train size: %d; dev size %d; test size: %d" % (
        len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


if __name__ == "__main__":
    read_ner("../ner_data/drugs/train.txt", "drugs", "wnut")

<filename>src/dataloader.py
import torch
import torch.nn as nn
import torch.utils.data as data
from torch.utils.data import DataLoader

import math
from tqdm import tqdm
import random
import logging

logger = logging.getLogger()

# from transformers import BertTokenizer
# bert_tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
from src.config import get_params

params = get_params()
from transformers import AutoTokenizer

auto_tokenizer = AutoTokenizer.from_pretrained(params.model_name)
pad_token_label_id = nn.CrossEntropyLoss().ignore_index


# domain2labels = {"politics": politics_labels, "science": science_labels, "music": music_labels, "literature": literature_labels, "ai": ai_labels }
# "drugs":conll_to_drugs, "single_drugs":single_drugs, "conll_to_drugs":conll_to_drugs, "wnut_to_drugs":wnut_to_drugs, "btc_to_drugs":btc_to_drugs

def remove_duplicates_with_order(seq):
    seen = set()
    seen_add = seen.add
    return [x for x in seq if not (x in seen or seen_add(x))]

def get_domain2labels(tgt_dm, src_dm):
    politics_labels = ['O', 'B-country', 'B-politician', 'I-politician', 'B-election', 'I-election', 'B-person',
                       'I-person', 'B-organisation', 'I-organisation', 'B-location', 'B-misc', 'I-location',
                       'I-country', 'I-misc', 'B-politicalparty', 'I-politicalparty', 'B-event', 'I-event']
    science_labels = ['O', 'B-scientist', 'I-scientist', 'B-person', 'I-person', 'B-university', 'I-university',
                      'B-organisation', 'I-organisation', 'B-country', 'I-country', 'B-location', 'I-location',
                      'B-discipline', 'I-discipline', 'B-enzyme', 'I-enzyme', 'B-protein', 'I-protein',
                      'B-chemicalelement', 'I-chemicalelement', 'B-chemicalcompound', 'I-chemicalcompound',
                      'B-astronomicalobject', 'I-astronomicalobject', 'B-academicjournal', 'I-academicjournal',
                      'B-event', 'I-event', 'B-theory', 'I-theory', 'B-award', 'I-award', 'B-misc', 'I-misc']
    music_labels = ['O', 'B-musicgenre', 'I-musicgenre', 'B-song', 'I-song', 'B-band', 'I-band', 'B-album', 'I-album',
                    'B-musicalartist', 'I-musicalartist', 'B-musicalinstrument', 'I-musicalinstrument', 'B-award',
                    'I-award', 'B-event', 'I-event', 'B-country', 'I-country', 'B-location', 'I-location',
                    'B-organisation', 'I-organisation', 'B-person', 'I-person', 'B-misc', 'I-misc']
    literature_labels = ["O", "B-book", "I-book", "B-writer", "I-writer", "B-award", "I-award", "B-poem", "I-poem",
                         "B-event", "I-event", "B-magazine", "I-magazine", "B-literarygenre", "I-literarygenre",
                         'B-country', 'I-country', "B-person", "I-person", "B-location", "I-location", 'B-organisation',
                         'I-organisation', 'B-misc', 'I-misc']
    ai_labels = ["O", "B-field", "I-field", "B-task", "I-task", "B-product", "I-product", "B-algorithm", "I-algorithm",
                 "B-researcher", "I-researcher", "B-metrics", "I-metrics", "B-programlang", "I-programlang",
                 "B-conference", "I-conference", "B-university", "I-university", "B-country", "I-country", "B-person",
                 "I-person", "B-organisation", "I-organisation", "B-location", "I-location", "B-misc", "I-misc"]

    single_wnut = ["O", "B-creative-work", "I-creative-work", "B-person", "I-person", "B-corporation", "I-corporation",
                   "B-location", "I-location", "B-group", "I-group", "B-product", "I-product"]  # 12 without "O"
    single_conll = ['O', 'B-organisation', 'I-organisation', 'B-misc', 'I-misc', 'B-person', 'I-person', 'B-location',
                    'I-location']  # 8 without "O"
    single_btc = ["O", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG", "I-ORG"]  # 6 without "O"
    single_drugs = ["O", "B-Drug", "I-Drug"]  # 2 without "O"

    conll_to_drugs = ['O', 'B-organisation', 'I-organisation', 'B-misc', 'I-misc', 'B-person', 'I-person', 'B-location',
                      'I-location', 'B-Drug', 'I-Drug']
    wnut_to_drugs = ["O", "B-creative-work", "I-creative-work", "B-person", "I-person", "B-corporation",
                     "I-corporation", "B-location", "I-location", "B-product", "I-product", "B-Drug", "I-Drug"]
    btc_to_drugs = ["O", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG", "I-ORG", "B-Drug", "I-Drug"]
    wnut_to_btc = ["O", "B-creative-work", "I-creative-work", "B-person", "I-person", "B-corporation", "I-corporation",
                   "B-location", "I-location", "B-product", "I-product", "B-PER", "I-PER", "B-LOC", "I-LOC", "B-ORG",
                   "I-ORG", "B-Drug", "I-Drug"]

    domain2labels = {"politics": politics_labels, "science": science_labels, "music": music_labels,
                     "literature": literature_labels, "ai": ai_labels}
    custom_domains = {'wnut': single_wnut, 'nutot': single_wnut, 'btc': single_btc, 'drugs': single_drugs,
                      'wiki': single_drugs, 'conll': single_conll}

    # create custom entity class set
    if tgt_dm not in domain2labels:
        if tgt_dm in custom_domains and src_dm in custom_domains:
            domain2labels[tgt_dm] = custom_domains[tgt_dm]
            domain2labels[tgt_dm].extend(custom_domains[src_dm])
            domain2labels[tgt_dm] = remove_duplicates_with_order(domain2labels[tgt_dm])  # remove duplicate entity classes
        else:
            raise Exception("Target or Source Domain not in Custom Domain Set! Fix that")
    elif params.tgt_dm not in domain2labels and params.src_dm in custom_domains and src_dm != 'conll':
        raise Exception("If you want to use a standard domain with custom source domain --> you gotta configure this.")

    return domain2labels

def split_text(token_list, label_list, cut_strat):
    dot_enc = auto_tokenizer.convert_tokens_to_ids('.')
    splitted_token_lists, splitted_label_lists = [], []
    amt_cutoff=0
    amt_total_splits=1


    if cut_strat != 1: # if cut_strat is 1 than directly cutoff long texts.
        if cut_strat == 2:
            start_at_indx = 400
            cut_until = 500
        elif cut_strat == 3:
            start_at_indx = 0#Cut into sentences
            cut_until = 0
        elif cut_strat == 4: #Ignore longer texts
            return splitted_token_lists, splitted_label_lists, len(token_list), amt_total_splits
        else:
            print('Cut Strategy: ' , str(cut_strat))
            raise ValueError
        while len(token_list) > cut_until and start_at_indx >= 0:
            try:
                cut_at = token_list.index(dot_enc, start_at_indx, 500)
                splitted_token_lists.append(token_list[:cut_at])
                splitted_label_lists.append(label_list[:cut_at])
                if cut_at < len(token_list):
                    token_list = token_list[cut_at+1:]
                    label_list = label_list[cut_at+1:]
                    amt_total_splits+=1
                else:
                    token_list = []
                    label_list = []

                if cut_strat == 2:#Reset starting point for next possible split
                    start_at_indx = 400
            except ValueError:
                logger.info("Split not found enlarging search area")
                start_at_indx = start_at_indx - 100 #If no dot is in that area --> make search are bigger or continue

    if len(token_list) > 500:
        amt_cutoff = len(token_list) - 500
        splitted_token_lists.append(token_list[:500])
        splitted_label_lists.append(label_list[:500])
    else:
        splitted_token_lists.append(token_list)
        splitted_label_lists.append(label_list)

    return splitted_token_lists, splitted_label_lists, amt_cutoff, amt_total_splits




def read_ner(datapath, tgt_dm, src_dm):
    inputs, labels = [], []
    domain2labels = get_domain2labels(tgt_dm, src_dm)
    amt_cut_off=0
    amt_tokens=0
    amt_splits=0

    with open(datapath, "r") as fr:
        token_list, label_list  = [], []

        for i, line in enumerate(fr):
            line = line.strip()
            if line == "":
                if len(token_list) > 0:
                    assert len(token_list) == len(label_list)

                    if len(token_list) > 500 or params.cut_strategy==3:
                        splitted_token_lists, splitted_label_lists, recent_cutoff, rec_splits = split_text(token_list, label_list, params.cut_strategy)
                        amt_cut_off += recent_cutoff
                        amt_splits += rec_splits



                        for cut_tokens in splitted_token_lists:
                            inputs.append([auto_tokenizer.cls_token_id] + cut_tokens + [auto_tokenizer.sep_token_id])
                        for cut_labels in splitted_label_lists:
                            labels.append([pad_token_label_id] + cut_labels + [pad_token_label_id])
                    else:
                        amt_tokens += len(token_list)
                        inputs.append([auto_tokenizer.cls_token_id] + token_list + [auto_tokenizer.sep_token_id])
                        labels.append([pad_token_label_id] + label_list + [pad_token_label_id])

                token_list, label_list = [], []
                continue

            splits = line.split("\t")

            if len(splits) != 2:
                print('i:%d, text:%s'%(i, line))

            token = splits[0]
            #Check if its wnut, since this dataset has multiple labels on the testset..
            if ( src_dm != 'wnut' and src_dm != 'nutot' and tgt_dm != 'wnut'  and tgt_dm != 'nutot' ) or "," not in splits[1]:
                label = splits[1]
            else: # If its WNUT and it has commas in the label --> prioritize product, but acutally doesn't matter much.#TODO check if I should build this in properly
                possible_labels = splits[1].split(',')
                if 'B-product' in possible_labels:
                    label = 'B-product'
                elif 'I-product' in possible_labels:
                    label = 'I-product'
                else:
                    label = possible_labels[0]


            subs_ = auto_tokenizer.tokenize(token)
            if len(subs_) > 0:
                label_list.extend([domain2labels[tgt_dm].index(label)] + [pad_token_label_id] * (len(subs_) - 1))
                token_list.extend(auto_tokenizer.convert_tokens_to_ids(subs_))
            else:
                print("length of subwords for %s is zero; its label is %s" % (token, label))

    logger.info('Read total of %d tokens and cutoff %d tokens for file %s and split/cut at least into %d times'%(amt_tokens, amt_cut_off, datapath, amt_splits))
    return inputs, labels


def read_ner_for_bilstm(datapath, tgt_dm, vocab):
    domain2labels = get_domain2labels(tgt_dm, 'Error')
    inputs, labels = [], []
    with open(datapath, "r") as fr:
        token_list, label_list = [], []
        for i, line in enumerate(fr):
            line = line.strip()
            if line == "":
                if len(token_list) > 0:
                    assert len(token_list) == len(label_list)
                    inputs.append(token_list)
                    labels.append(label_list)

                token_list, label_list = [], []
                continue

            splits = line.split("\t")
            token = splits[0]
            label = splits[1]

            token_list.append(vocab.word2index[token])
            label_list.append(domain2labels[tgt_dm].index(label))

    return inputs, labels


class Dataset(data.Dataset):
    def __init__(self, inputs, labels):
        self.X = inputs
        self.y = labels

    def __getitem__(self, index):
        return self.X[index], self.y[index]

    def __len__(self):
        return len(self.X)


PAD_INDEX = 0


class Vocab():
    def __init__(self):
        self.word2index = {"PAD": PAD_INDEX}
        self.index2word = {PAD_INDEX: "PAD"}
        self.n_words = 1

    def index_words(self, word_list):
        for word in word_list:
            if word not in self.word2index:
                self.word2index[word] = self.n_words
                self.index2word[self.n_words] = word
                self.n_words += 1


def get_vocab(path):
    vocabulary = []
    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            vocabulary.append(line)
    return vocabulary


def collate_fn(data):
    X, y = zip(*data)
    lengths = [len(bs_x) for bs_x in X]
    max_lengths = max(lengths)
    padded_seqs = torch.LongTensor(len(X), max_lengths).fill_(auto_tokenizer.pad_token_id)
    padded_y = torch.LongTensor(len(X), max_lengths).fill_(pad_token_label_id)
    for i, (seq, y_) in enumerate(zip(X, y)):
        length = lengths[i]
        padded_seqs[i, :length] = torch.LongTensor(seq)
        padded_y[i, :length] = torch.LongTensor(y_)

    return padded_seqs, padded_y


def collate_fn_for_bilstm(data):
    X, y = zip(*data)
    lengths = [len(bs_x) for bs_x in X]
    max_lengths = max(lengths)
    padded_seqs = torch.LongTensor(len(X), max_lengths).fill_(PAD_INDEX)
    for i, seq in enumerate(X):
        length = lengths[i]
        padded_seqs[i, :length] = torch.LongTensor(seq)

    lengths = torch.LongTensor(lengths)
    return padded_seqs, lengths, y


def get_dataloader_for_bilstmtagger(params):
    vocab_src = get_vocab("ner_data/conll2003/vocab.txt")
    vocab_tgt = get_vocab("ner_data/%s/vocab.txt" % params.tgt_dm)
    vocab = Vocab()
    vocab.index_words(vocab_src)
    vocab.index_words(vocab_tgt)

    logger.info("Load training set data ...")
    conll_inputs_train, conll_labels_train = read_ner_for_bilstm("ner_data/conll2003/train.txt", params.tgt_dm, vocab)
    inputs_train, labels_train = read_ner_for_bilstm("ner_data/%s/train.txt" % params.tgt_dm, params.tgt_dm, vocab)
    inputs_train = inputs_train * 10 + conll_inputs_train
    labels_train = labels_train * 10 + conll_labels_train

    logger.info("Load dev set data ...")
    inputs_dev, labels_dev = read_ner_for_bilstm("ner_data/%s/dev.txt" % params.tgt_dm, params.tgt_dm, vocab)

    logger.info("Load test set data ...")
    inputs_test, labels_test = read_ner_for_bilstm("ner_data/%s/test.txt" % params.tgt_dm, params.tgt_dm, vocab)

    logger.info("train size: %d; dev size %d; test size: %d;" % (len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=params.batch_size, shuffle=True,
                                  collate_fn=collate_fn_for_bilstm)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=params.batch_size, shuffle=False,
                                collate_fn=collate_fn_for_bilstm)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=params.batch_size, shuffle=False,
                                 collate_fn=collate_fn_for_bilstm)

    return dataloader_train, dataloader_dev, dataloader_test, vocab


def load_corpus(tgt_dm):
    print("Loading corpus ...")
    data_path = "enwiki_corpus/%s_removebracket.tok" % tgt_dm
    sent_list = []
    with open(data_path, "r") as fr:
        for i, line in tqdm(enumerate(fr)):
            line = line.strip()
            sent_list.append(line)
    return sent_list


def get_dataloader(params):
    logger.info("Load training set data")
    inputs_train, labels_train = read_ner("ner_data/%s/train.txt" % params.tgt_dm, params.tgt_dm, params.src_dm)
    domain2labels = get_domain2labels(params.tgt_dm, params.src_dm)

    if params.n_samples != -1:
        logger.info("Few-shot on %d samples" % params.n_samples)
        inputs_train = inputs_train[:params.n_samples]
        labels_train = labels_train[:params.n_samples]
    logger.info("Load development set data")
    inputs_dev, labels_dev = read_ner("ner_data/%s/dev.txt" % params.tgt_dm, params.tgt_dm, params.src_dm)
    logger.info("Load test set data")
    inputs_test, labels_test = read_ner("ner_data/%s/test.txt" % params.tgt_dm, params.tgt_dm, params.src_dm)

    logger.info("label distribution for training set")
    label_distri_train = {}
    count_tok_train = 0
    for label_seq in labels_train:
        for label in label_seq:
            if label != pad_token_label_id:
                label_name = domain2labels[params.tgt_dm][label]
                if "B" in label_name:
                    count_tok_train += 1
                    label_name = label_name.split("-")[1]
                    if label_name not in label_distri_train:
                        label_distri_train[label_name] = 1
                    else:
                        label_distri_train[label_name] += 1
    print(label_distri_train)
    for key in label_distri_train:
        label_distri_train[key] /= count_tok_train
    logger.info(label_distri_train)

    logger.info("label distribution for dev set")
    label_distri_dev = {}
    count_tok_test = 0
    for label_seq in labels_dev:
        for label in label_seq:
            if label != pad_token_label_id:
                label_name = domain2labels[params.tgt_dm][label]
                if "B" in label_name:
                    count_tok_test += 1
                    label_name = label_name.split("-")[1]
                    if label_name not in label_distri_dev:
                        label_distri_dev[label_name] = 1
                    else:
                        label_distri_dev[label_name] += 1
    print(label_distri_dev)
    for key in label_distri_dev:
        label_distri_dev[key] /= count_tok_test
    logger.info(label_distri_dev)

    logger.info("label distribution for test set")
    label_distri_test = {}
    count_tok_test = 0
    for label_seq in labels_test:
        for label in label_seq:
            if label != pad_token_label_id:
                label_name = domain2labels[params.tgt_dm][label]
                if "B" in label_name:
                    count_tok_test += 1
                    label_name = label_name.split("-")[1]
                    if label_name not in label_distri_test:
                        label_distri_test[label_name] = 1
                    else:
                        label_distri_test[label_name] += 1
    print(label_distri_test)
    for key in label_distri_test:
        label_distri_test[key] /= count_tok_test
    logger.info(label_distri_test)

    if params.conll and params.joint: #No worries, this is only done when training is executed in a joint manner. Furthermore our dataset are nearly equal so balancing is not so important.
        conll_inputs_train, conll_labels_train = read_ner("ner_data/conll2003/train.txt", params.tgt_dm, params.src_dm)
        inputs_train = inputs_train * 50  # augment the target domain data to balance the source and target domain data
        labels_train = labels_train * 50
        inputs_train = inputs_train + conll_inputs_train
        labels_train = labels_train + conll_labels_train

    logger.info("train size: %d; dev size %d; test size: %d;" % (len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=params.batch_size, shuffle=True,
                                  collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=params.batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=params.batch_size, shuffle=False,
                                 collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_conll2003_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/conll2003/train.txt", params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/conll2003/dev.txt", params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/conll2003/test.txt", params.tgt_dm, params.src_dm)

    logger.info("conll2003 dataset: train size: %d; dev size %d; test size: %d" % (
    len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_btc_dataloader(batch_size, params):
    # According to https://github.com/GateNLP/broad_twitter_corpus
    # Training: everything else
    inputs_train1, labels_train1 = read_ner("ner_data/btc/a.conll", params.tgt_dm, params.src_dm)
    inputs_train2, labels_train2 = read_ner("ner_data/btc/b.conll", params.tgt_dm, params.src_dm)
    inputs_train3, labels_train3 = read_ner("ner_data/btc/e.conll", params.tgt_dm, params.src_dm)
    inputs_train4, labels_train4 = read_ner("ner_data/btc/g.conll", params.tgt_dm, params.src_dm)
    inputs_train5, labels_train5 = read_ner("ner_data/btc/h.conll", params.tgt_dm, params.src_dm)

    # Development: second half of Section H
    inputs_dev = inputs_train5[math.floor( len(inputs_train5) / 2):]
    labels_dev = labels_train5[math.floor( len(labels_train5) / 2):]
    inputs_train5 = inputs_train5[:math.floor( len(inputs_train5) / 2)]
    labels_train5 = labels_train5[:math.floor( len(labels_train5) / 2)]

    # Test: Section F
    inputs_test, labels_test = read_ner("ner_data/btc/f.conll", params.tgt_dm, params.src_dm)

    # Aggregate Train sets
    inputs_train = inputs_train1 + inputs_train2 + inputs_train3 + inputs_train4 + inputs_train5
    labels_train = labels_train1 + labels_train2 + labels_train3 + labels_train4 + labels_train5

    logger.info("Broad Twitter Corpus dataset: train size: %d; dev size %d; test size: %d" % (
        len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_wnut_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/wnut/train.txt", params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/wnut/dev.txt", params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/wnut/test.txt", params.tgt_dm, params.src_dm)

    logger.info("W-NUT 2017 dataset: train size: %d; dev size %d; test size: %d" % (
        len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_generic_source_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/%s/train.txt" % params.src_dm, params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/%s/dev.txt" % params.src_dm, params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/%s/test.txt" % params.src_dm, params.tgt_dm, params.src_dm)

    logger.info("%s dataset: train size: %d; dev size %d; test size: %d" % (params.src_dm,
                                                                            len(inputs_train), len(inputs_dev),
                                                                            len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


def get_wiki_dataloader(batch_size, params):
    inputs_train, labels_train = read_ner("ner_data/wiki/train.txt", params.tgt_dm, params.src_dm)
    inputs_dev, labels_dev = read_ner("ner_data/wiki/dev.txt", params.tgt_dm, params.src_dm)
    inputs_test, labels_test = read_ner("ner_data/wiki/test.txt", params.tgt_dm, params.src_dm)

    logger.info("Wikipedia distantly supervised dataset: train size: %d; dev size %d; test size: %d" % (
        len(inputs_train), len(inputs_dev), len(inputs_test)))

    dataset_train = Dataset(inputs_train, labels_train)
    dataset_dev = Dataset(inputs_dev, labels_dev)
    dataset_test = Dataset(inputs_test, labels_test)

    dataloader_train = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    dataloader_dev = DataLoader(dataset=dataset_dev, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
    dataloader_test = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

    return dataloader_train, dataloader_dev, dataloader_test


if __name__ == "__main__":
    read_ner("../ner_data/drugs/train.txt", "drugs", "wnut")

# from transformers import BertTokenizer # bert_tokenizer = BertTokenizer.from_pretrained('bert-base-cased') # domain2labels = {"politics": politics_labels, "science": science_labels, "music": music_labels, "literature": literature_labels, "ai": ai_labels } # "drugs":conll_to_drugs, "single_drugs":single_drugs, "conll_to_drugs":conll_to_drugs, "wnut_to_drugs":wnut_to_drugs, "btc_to_drugs":btc_to_drugs # 12 without "O" # 8 without "O" # 6 without "O" # 2 without "O" # create custom entity class set # remove duplicate entity classes # if cut_strat is 1 than directly cutoff long texts. #Cut into sentences #Ignore longer texts #Reset starting point for next possible split #If no dot is in that area --> make search are bigger or continue #Check if its wnut, since this dataset has multiple labels on the testset.. # If its WNUT and it has commas in the label --> prioritize product, but acutally doesn't matter much.#TODO check if I should build this in properly #No worries, this is only done when training is executed in a joint manner. Furthermore our dataset are nearly equal so balancing is not so important. # augment the target domain data to balance the source and target domain data # According to https://github.com/GateNLP/broad_twitter_corpus # Training: everything else # Development: second half of Section H # Test: Section F # Aggregate Train sets