#!/usr/bin/env python
# -*- coding: utf-8 -*-
######################################################################
#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
######################################################################
"""
File: data_provider.py
"""

import re
import os
import types
import csv
import random
import numpy as np


VOC_DICT = {}


def load_dict(vocab_dict): 
    """
    load vocabulary dict
    """
    idx = 0
    for line in open(vocab_dict): 
        line = line.strip()
        VOC_DICT[line] = idx
        idx += 1
    return VOC_DICT


def prepare_batch_data(insts,
                       task_name,
                       max_len=128,
                       return_attn_bias=True,
                       return_max_len=True,
                       return_num_token=False):
    """
    generate self attention mask, [shape: batch_size *  max_len * max_len]
    """
    batch_context_ids = [inst[0] for inst in insts]
    batch_context_pos_ids = [inst[1] for inst in insts]
    batch_segment_ids = [inst[2] for inst in insts]
    batch_label_ids = [[inst[3]] for inst in insts]
    labels_list = batch_label_ids

    context_id, next_sent_context_index, context_attn_bias = \
            pad_batch_data(batch_context_ids, pad_idx=0, max_len=max_len, \
            return_next_sent_pos=True, return_attn_bias=True)

    context_pos_id = pad_batch_data(
        batch_context_pos_ids, pad_idx=0, max_len=max_len, return_pos=False, return_attn_bias=False)

    context_segment_id = pad_batch_data(
        batch_segment_ids, pad_idx=0, max_len=max_len, return_pos=False, return_attn_bias=False)

    if 'kn' in task_name: 
        batch_kn_ids = [inst[4] for inst in insts]
        kn_id = pad_bath_kn_data(batch_kn_ids, pad_idx=0, max_len=max_len)

        batch_goal_ids = [inst[5] for inst in insts]
        goal_id = pad_bath_kn_data(batch_goal_ids, pad_idx=0, max_len=max_len)

    out_list = []
    for i in range(len(insts)): 
        if 'kn' in task_name: 
            out = [context_id[i], context_pos_id[i], context_segment_id[i], context_attn_bias[i], \
                    kn_id[i], goal_id[i], labels_list[i], next_sent_context_index[i]]
        else: 
            out = [context_id[i], context_pos_id[i], context_segment_id[i], \
                    context_attn_bias[i], labels_list[i], next_sent_context_index[i]]
        out_list.append(out)
    return out_list


def pad_bath_kn_data(insts,
                    pad_idx=0,
                    max_len=128):
    """pad_bath_kn_data"""
    kn_list = []
    for inst in insts: 
        inst = inst[0: min(max_len, len(inst))]
        kn_list.append(inst)
    return kn_list


def pad_batch_data(insts,
                   pad_idx=0,
                   max_len=128,
                   return_pos=False,
                   return_next_sent_pos=False,
                   return_attn_bias=False,
                   return_max_len=False,
                   return_num_token=False):
    """
    Pad the instances to the max sequence length in batch, and generate the
    corresponding position data and attention bias.
    """
    return_list = []

    inst_data = np.array(
        [inst + list([pad_idx] * (max_len - len(inst))) for inst in insts])
    return_list += [inst_data.astype("int64").reshape([-1, max_len, 1])]

    if return_next_sent_pos:
        batch_size = inst_data.shape[0]
        max_seq_len = inst_data.shape[1]
        next_sent_index = np.array(
            range(0, batch_size * max_seq_len, max_seq_len)).astype(
                "int64").reshape(-1, 1)
        return_list += [next_sent_index]

    if return_pos:
        inst_pos = np.array([
            list(range(0, len(inst))) + [pad_idx] * (max_len - len(inst))
            for inst in insts])
        return_list += [inst_pos.astype("int64").reshape([-1, max_len, 1])]

    if return_attn_bias:
        slf_attn_bias_data = np.array([[0] * len(inst) + [-1e9] *
                                       (max_len - len(inst)) for inst in insts])
        slf_attn_bias_data = np.tile(
            slf_attn_bias_data.reshape([-1, 1, max_len]), [1, max_len, 1])
        return_list += [slf_attn_bias_data.astype("float32")]

    if return_max_len:
        return_list += [max_len]

    if return_num_token:
        num_token = 0
        for inst in insts:
            num_token += len(inst)
        return_list += [num_token]

    return return_list if len(return_list) > 1 else return_list[0]


def preprocessing_for_one_line(line, labels, task_name, max_seq_len=256):
    """
    process text to model inputs
    """
    line = line.rstrip('\n').split('\t')
    label_text = line[0]
    context_text = line[1]
    response_text = line[2]
    if 'kn' in task_name:
        kn_text = "%s [SEP] %s" % (line[3], line[4])
    else:
        kn_text = None

    example = InputExample(guid=0, \
                           context_text=context_text, \
                           response_text=response_text, \
                           kn_text=kn_text, \
                           goal_text = kn_text, \
                           label_text=label_text)

    feature = convert_single_example(0, example, labels, max_seq_len)

    instance = [feature.context_ids, feature.context_pos_ids, \
                feature.segment_ids, feature.label_ids, feature.kn_ids]

    batch_data = prepare_batch_data([instance],
                                    task_name,
                                    max_len=max_seq_len,
                                    return_attn_bias=True,
                                    return_max_len=False,
                                    return_num_token=False)

    return batch_data


class DataProcessor(object):
    """Base class for data converters for sequence classification data sets."""
    def __init__(self, data_dir, task_name, vocab_path, max_seq_len, do_lower_case):
        self.data_dir = data_dir
        self.max_seq_len = max_seq_len
        self.task_name = task_name

        self.current_train_example = -1
        self.num_examples = {'train': -1, 'dev': -1, 'test': -1}
        self.current_train_epoch = -1
        VOC_DICT = load_dict(vocab_path)

    def get_train_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

    def get_dev_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

    def get_test_examples(self, data_dir):
        """Gets a collection of `InputExample`s for prediction."""
        raise NotImplementedError()

    @classmethod
    def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

    def convert_example(self, index, example, labels, max_seq_len):
        """Converts a single `InputExample` into a single `InputFeatures`."""
        feature = convert_single_example(index, example, labels, max_seq_len)
        return feature

    def generate_batch_data(self,
                            batch_data,
                            voc_size=-1,
                            mask_id=-1,
                            return_attn_bias=True,
                            return_max_len=False,
                            return_num_token=False):
        """generate_batch_data"""
        return prepare_batch_data(
            batch_data,
            self.task_name,
            self.max_seq_len,
            return_attn_bias=True,
            return_max_len=False,
            return_num_token=False)

    @classmethod
    def _read_data(cls, input_file):
        """Reads a tab separated value file."""
        with open(input_file, "r") as f:
            lines = []
            for line in f:
                line = line.rstrip('\n').split('\t')
                lines.append(line)
            return lines

    def get_num_examples(self, phase):
        """Get number of examples for train, dev or test."""
        if phase not in ['train', 'dev', 'test']:
            raise ValueError("Unknown phase, which should be in ['train', 'dev', 'test'].")
        return self.num_examples[phase]

    def get_train_progress(self):
        """Gets progress for training phase."""
        return self.current_train_example, self.current_train_epoch

    def data_generator(self,
                       batch_size,
                       phase='train',
                       epoch=1,
                       shuffle=False):
        """
        Generate data for train, dev or test.
        """
        if phase == 'train':
            examples = self.get_train_examples(self.data_dir)
            self.num_examples['train'] = len(examples)
        elif phase == 'dev':
            examples = self.get_dev_examples(self.data_dir)
            self.num_examples['dev'] = len(examples)
        elif phase == 'test':
            examples = self.get_test_examples(self.data_dir)
            self.num_examples['test'] = len(examples)
        else:
            raise ValueError("Unknown phase, which should be in ['train', 'dev', 'test'].")

        def instance_reader():
            """instance_reader"""
            for epoch_index in range(epoch):
                if shuffle:
                    random.shuffle(examples)
                if phase == 'train':
                    self.current_train_epoch = epoch_index
                for (index, example) in enumerate(examples):
                    if phase == 'train':
                        self.current_train_example = index + 1
                    feature = self.convert_example(
                        index, example, self.get_labels(), self.max_seq_len)
                    if 'kn' in self.task_name: 
                        instance = [feature.context_ids, feature.context_pos_ids, \
                                feature.segment_ids, feature.label_ids, feature.kn_ids, feature.goal_ids]
                    else: 
                        instance = [feature.context_ids, feature.context_pos_ids, \
                                feature.segment_ids, feature.label_ids]
                    yield instance

        def batch_reader(reader, batch_size):
            """batch_reader"""
            batch = []
            for instance in reader():
                if len(batch) < batch_size: 
                    batch.append(instance)
                else: 
                    yield batch
                    batch = [instance]

            if len(batch) > 0:
                yield batch

        def wrapper():
            """wrapper"""
            for batch_data in batch_reader(instance_reader, batch_size):
                batch_data = self.generate_batch_data(
                    batch_data,
                    voc_size=-1,
                    mask_id=-1,
                    return_attn_bias=True,
                    return_max_len=False,
                    return_num_token=False)
                yield batch_data

        return wrapper


class InputExample(object):
    """A single training/test example"""

    def __init__(self, guid, context_text, response_text, kn_text, goal_text, label_text):
        self.guid = guid
        self.context_text = context_text
        self.response_text = response_text
        self.kn_text = kn_text
        self.goal_text = goal_text
        self.label_text = label_text


class InputFeatures(object): 
    """input features datas"""
    def __init__(self, context_ids, context_pos_ids, segment_ids, kn_ids, goal_ids, label_ids): 
        self.context_ids = context_ids
        self.context_pos_ids = context_pos_ids
        self.segment_ids = segment_ids
        self.kn_ids = kn_ids
        self.goal_ids = goal_ids
        self.label_ids = label_ids


class MatchProcessor(DataProcessor):
    """Processor for the Match data set (GLUE version)."""

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_data(os.path.join(data_dir, "train.txt")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_data(os.path.join(data_dir, "dev.txt")), "dev")

    def get_test_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_data(os.path.join(data_dir, "test.txt")), "test")

    @classmethod
    def get_labels(self):
        """See base class."""
        return ["0", "1"]

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)
            context_text = line[1]
            label_text = line[0]
            response_text = line[2]
            if 'kn' in self.task_name: 
                kn_text = line[4]
                goal_text = line[3]
            else: 
                kn_text = None
                goal_text = None
            examples.append(
                InputExample(
                    guid=guid, context_text=context_text, response_text=response_text, \
                            kn_text=kn_text, goal_text=goal_text, label_text=label_text))
        return examples


def convert_tokens_to_ids(tokens): 
    """
    convert input ids
    """
    ids = []
    for token in tokens: 
        if token in VOC_DICT: 
            ids.append(VOC_DICT[token])
        else: 
            ids.append(VOC_DICT['[UNK]'])
    return ids


def convert_single_example(ex_index, example, label_list, max_seq_length):
    """Converts a single `InputExample` into a single `InputFeatures`."""
    label_map = {}
    for (i, label) in enumerate(label_list):
        label_map[label] = i
    if example.context_text: 
        tokens_context = example.context_text
        tokens_context = tokens_context.split()
    else: 
        tokens_context = []

    if example.response_text: 
        tokens_response = example.response_text
        tokens_response = tokens_response.split()
    else: 
        tokens_response = []

    if example.kn_text:
        tokens_kn = example.kn_text
        tokens_kn = tokens_kn.split()
        tokens_kn = tokens_kn[0: min(len(tokens_kn), max_seq_length)]
    else: 
        tokens_kn = []

    if example.goal_text:
        tokens_goal = example.goal_text
        tokens_goal = tokens_goal.split()
        tokens_goal = tokens_goal[0: min(len(tokens_goal), max_seq_length)]
    else: 
        tokens_goal = []

    tokens_response = tokens_response[0: min(50, len(tokens_response))]
    if len(tokens_context) > max_seq_length - len(tokens_response) - 3: 
        tokens_context = tokens_context[len(tokens_context) \
                + len(tokens_response) - max_seq_length + 3:]

    context_tokens = []
    segment_ids = []

    context_tokens.append("[CLS]")
    segment_ids.append(0)
    context_tokens.extend(tokens_context)
    segment_ids.extend([0] * len(tokens_context))
    context_tokens.append("[SEP]")
    segment_ids.append(0)

    context_tokens.extend(tokens_response)
    segment_ids.extend([1] * len(tokens_response))
    context_tokens.append("[SEP]")
    segment_ids.append(1)

    context_ids = convert_tokens_to_ids(context_tokens)
    context_pos_ids = list(range(len(context_ids)))
    label_ids = label_map[example.label_text]
    if tokens_kn: 
        kn_ids = convert_tokens_to_ids(tokens_kn)
    else: 
        kn_ids = []

    if tokens_goal: 
        goal_ids = convert_tokens_to_ids(tokens_goal)
    else: 
        goal_ids = []

    feature = InputFeatures(
        context_ids=context_ids,
        context_pos_ids=context_pos_ids,
        segment_ids=segment_ids,
        kn_ids = kn_ids,
        goal_ids = goal_ids,
        label_ids=label_ids)
    #if ex_index < 5:
    #    print("*** Example ***")
    #    print("guid: %s" % (example.guid))
    #    print("context tokens: %s" % " ".join(context_tokens))
    #    print("context_ids: %s" % " ".join([str(x) for x in context_ids]))
    #    print("context_pos_ids: %s" % " ".join([str(x) for x in context_pos_ids]))
    #    print("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
    #    print("kn_ids: %s" % " ".join([str(x) for x in kn_ids]))
    #    print("label: %s (id = %d)" % (example.label_text, label_ids))
    return feature


#!/usr/bin/env python
# -*- coding: utf-8 -*-
######################################################################
#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
######################################################################
"""
File: data_provider.py
"""

import re
import os
import types
import csv
import random
import numpy as np


VOC_DICT = {}


def load_dict(vocab_dict): 
    """
    load vocabulary dict
    """
    idx = 0
    for line in open(vocab_dict): 
        line = line.strip()
        VOC_DICT[line] = idx
        idx += 1
    return VOC_DICT


def prepare_batch_data(insts,
                       task_name,
                       max_len=128,
                       return_attn_bias=True,
                       return_max_len=True,
                       return_num_token=False):
    """
    generate self attention mask, [shape: batch_size *  max_len * max_len]
    """
    batch_context_ids = [inst[0] for inst in insts]
    batch_context_pos_ids = [inst[1] for inst in insts]
    batch_segment_ids = [inst[2] for inst in insts]
    batch_label_ids = [[inst[3]] for inst in insts]
    labels_list = batch_label_ids

    context_id, next_sent_context_index, context_attn_bias = \
            pad_batch_data(batch_context_ids, pad_idx=0, max_len=max_len, \
            return_next_sent_pos=True, return_attn_bias=True)

    context_pos_id = pad_batch_data(
        batch_context_pos_ids, pad_idx=0, max_len=max_len, return_pos=False, return_attn_bias=False)

    context_segment_id = pad_batch_data(
        batch_segment_ids, pad_idx=0, max_len=max_len, return_pos=False, return_attn_bias=False)

    if 'kn' in task_name: 
        batch_kn_ids = [inst[4] for inst in insts]
        kn_id = pad_bath_kn_data(batch_kn_ids, pad_idx=0, max_len=max_len)

        batch_goal_ids = [inst[5] for inst in insts]
        goal_id = pad_bath_kn_data(batch_goal_ids, pad_idx=0, max_len=max_len)

    out_list = []
    for i in range(len(insts)): 
        if 'kn' in task_name: 
            out = [context_id[i], context_pos_id[i], context_segment_id[i], context_attn_bias[i], \
                    kn_id[i], goal_id[i], labels_list[i], next_sent_context_index[i]]
        else: 
            out = [context_id[i], context_pos_id[i], context_segment_id[i], \
                    context_attn_bias[i], labels_list[i], next_sent_context_index[i]]
        out_list.append(out)
    return out_list


def pad_bath_kn_data(insts,
                    pad_idx=0,
                    max_len=128):
    """pad_bath_kn_data"""
    kn_list = []
    for inst in insts: 
        inst = inst[0: min(max_len, len(inst))]
        kn_list.append(inst)
    return kn_list


def pad_batch_data(insts,
                   pad_idx=0,
                   max_len=128,
                   return_pos=False,
                   return_next_sent_pos=False,
                   return_attn_bias=False,
                   return_max_len=False,
                   return_num_token=False):
    """
    Pad the instances to the max sequence length in batch, and generate the
    corresponding position data and attention bias.
    """
    return_list = []

    inst_data = np.array(
        [inst + list([pad_idx] * (max_len - len(inst))) for inst in insts])
    return_list += [inst_data.astype("int64").reshape([-1, max_len, 1])]

    if return_next_sent_pos:
        batch_size = inst_data.shape[0]
        max_seq_len = inst_data.shape[1]
        next_sent_index = np.array(
            range(0, batch_size * max_seq_len, max_seq_len)).astype(
                "int64").reshape(-1, 1)
        return_list += [next_sent_index]

    if return_pos:
        inst_pos = np.array([
            list(range(0, len(inst))) + [pad_idx] * (max_len - len(inst))
            for inst in insts])
        return_list += [inst_pos.astype("int64").reshape([-1, max_len, 1])]

    if return_attn_bias:
        slf_attn_bias_data = np.array([[0] * len(inst) + [-1e9] *
                                       (max_len - len(inst)) for inst in insts])
        slf_attn_bias_data = np.tile(
            slf_attn_bias_data.reshape([-1, 1, max_len]), [1, max_len, 1])
        return_list += [slf_attn_bias_data.astype("float32")]

    if return_max_len:
        return_list += [max_len]

    if return_num_token:
        num_token = 0
        for inst in insts:
            num_token += len(inst)
        return_list += [num_token]

    return return_list if len(return_list) > 1 else return_list[0]


def preprocessing_for_one_line(line, labels, task_name, max_seq_len=256):
    """
    process text to model inputs
    """
    line = line.rstrip('\n').split('\t')
    label_text = line[0]
    context_text = line[1]
    response_text = line[2]
    if 'kn' in task_name:
        kn_text = "%s [SEP] %s" % (line[3], line[4])
    else:
        kn_text = None

    example = InputExample(guid=0, \
                           context_text=context_text, \
                           response_text=response_text, \
                           kn_text=kn_text, \
                           goal_text = kn_text, \
                           label_text=label_text)

    feature = convert_single_example(0, example, labels, max_seq_len)

    instance = [feature.context_ids, feature.context_pos_ids, \
                feature.segment_ids, feature.label_ids, feature.kn_ids]

    batch_data = prepare_batch_data([instance],
                                    task_name,
                                    max_len=max_seq_len,
                                    return_attn_bias=True,
                                    return_max_len=False,
                                    return_num_token=False)

    return batch_data


class DataProcessor(object):
    """Base class for data converters for sequence classification data sets."""
    def __init__(self, data_dir, task_name, vocab_path, max_seq_len, do_lower_case):
        self.data_dir = data_dir
        self.max_seq_len = max_seq_len
        self.task_name = task_name

        self.current_train_example = -1
        self.num_examples = {'train': -1, 'dev': -1, 'test': -1}
        self.current_train_epoch = -1
        VOC_DICT = load_dict(vocab_path)

    def get_train_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

    def get_dev_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

    def get_test_examples(self, data_dir):
        """Gets a collection of `InputExample`s for prediction."""
        raise NotImplementedError()

    @classmethod
    def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

    def convert_example(self, index, example, labels, max_seq_len):
        """Converts a single `InputExample` into a single `InputFeatures`."""
        feature = convert_single_example(index, example, labels, max_seq_len)
        return feature

    def generate_batch_data(self,
                            batch_data,
                            voc_size=-1,
                            mask_id=-1,
                            return_attn_bias=True,
                            return_max_len=False,
                            return_num_token=False):
        """generate_batch_data"""
        return prepare_batch_data(
            batch_data,
            self.task_name,
            self.max_seq_len,
            return_attn_bias=True,
            return_max_len=False,
            return_num_token=False)

    @classmethod
    def _read_data(cls, input_file):
        """Reads a tab separated value file."""
        with open(input_file, "r") as f:
            lines = []
            for line in f:
                line = line.rstrip('\n').split('\t')
                lines.append(line)
            return lines

    def get_num_examples(self, phase):
        """Get number of examples for train, dev or test."""
        if phase not in ['train', 'dev', 'test']:
            raise ValueError("Unknown phase, which should be in ['train', 'dev', 'test'].")
        return self.num_examples[phase]

    def get_train_progress(self):
        """Gets progress for training phase."""
        return self.current_train_example, self.current_train_epoch

    def data_generator(self,
                       batch_size,
                       phase='train',
                       epoch=1,
                       shuffle=False):
        """
        Generate data for train, dev or test.
        """
        if phase == 'train':
            examples = self.get_train_examples(self.data_dir)
            self.num_examples['train'] = len(examples)
        elif phase == 'dev':
            examples = self.get_dev_examples(self.data_dir)
            self.num_examples['dev'] = len(examples)
        elif phase == 'test':
            examples = self.get_test_examples(self.data_dir)
            self.num_examples['test'] = len(examples)
        else:
            raise ValueError("Unknown phase, which should be in ['train', 'dev', 'test'].")

        def instance_reader():
            """instance_reader"""
            for epoch_index in range(epoch):
                if shuffle:
                    random.shuffle(examples)
                if phase == 'train':
                    self.current_train_epoch = epoch_index
                for (index, example) in enumerate(examples):
                    if phase == 'train':
                        self.current_train_example = index + 1
                    feature = self.convert_example(
                        index, example, self.get_labels(), self.max_seq_len)
                    if 'kn' in self.task_name: 
                        instance = [feature.context_ids, feature.context_pos_ids, \
                                feature.segment_ids, feature.label_ids, feature.kn_ids, feature.goal_ids]
                    else: 
                        instance = [feature.context_ids, feature.context_pos_ids, \
                                feature.segment_ids, feature.label_ids]
                    yield instance

        def batch_reader(reader, batch_size):
            """batch_reader"""
            batch = []
            for instance in reader():
                if len(batch) < batch_size: 
                    batch.append(instance)
                else: 
                    yield batch
                    batch = [instance]

            if len(batch) > 0:
                yield batch

        def wrapper():
            """wrapper"""
            for batch_data in batch_reader(instance_reader, batch_size):
                batch_data = self.generate_batch_data(
                    batch_data,
                    voc_size=-1,
                    mask_id=-1,
                    return_attn_bias=True,
                    return_max_len=False,
                    return_num_token=False)
                yield batch_data

        return wrapper


class InputExample(object):
    """A single training/test example"""

    def __init__(self, guid, context_text, response_text, kn_text, goal_text, label_text):
        self.guid = guid
        self.context_text = context_text
        self.response_text = response_text
        self.kn_text = kn_text
        self.goal_text = goal_text
        self.label_text = label_text


class InputFeatures(object): 
    """input features datas"""
    def __init__(self, context_ids, context_pos_ids, segment_ids, kn_ids, goal_ids, label_ids): 
        self.context_ids = context_ids
        self.context_pos_ids = context_pos_ids
        self.segment_ids = segment_ids
        self.kn_ids = kn_ids
        self.goal_ids = goal_ids
        self.label_ids = label_ids


class MatchProcessor(DataProcessor):
    """Processor for the Match data set (GLUE version)."""

    def get_train_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_data(os.path.join(data_dir, "train.txt")), "train")

    def get_dev_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_data(os.path.join(data_dir, "dev.txt")), "dev")

    def get_test_examples(self, data_dir):
        """See base class."""
        return self._create_examples(
            self._read_data(os.path.join(data_dir, "test.txt")), "test")

    @classmethod
    def get_labels(self):
        """See base class."""
        return ["0", "1"]

    def _create_examples(self, lines, set_type):
        """Creates examples for the training and dev sets."""
        examples = []
        for (i, line) in enumerate(lines):
            guid = "%s-%s" % (set_type, i)
            context_text = line[1]
            label_text = line[0]
            response_text = line[2]
            if 'kn' in self.task_name: 
                kn_text = line[4]
                goal_text = line[3]
            else: 
                kn_text = None
                goal_text = None
            examples.append(
                InputExample(
                    guid=guid, context_text=context_text, response_text=response_text, \
                            kn_text=kn_text, goal_text=goal_text, label_text=label_text))
        return examples


def convert_tokens_to_ids(tokens): 
    """
    convert input ids
    """
    ids = []
    for token in tokens: 
        if token in VOC_DICT: 
            ids.append(VOC_DICT[token])
        else: 
            ids.append(VOC_DICT['[UNK]'])
    return ids


def convert_single_example(ex_index, example, label_list, max_seq_length):
    """Converts a single `InputExample` into a single `InputFeatures`."""
    label_map = {}
    for (i, label) in enumerate(label_list):
        label_map[label] = i
    if example.context_text: 
        tokens_context = example.context_text
        tokens_context = tokens_context.split()
    else: 
        tokens_context = []

    if example.response_text: 
        tokens_response = example.response_text
        tokens_response = tokens_response.split()
    else: 
        tokens_response = []

    if example.kn_text:
        tokens_kn = example.kn_text
        tokens_kn = tokens_kn.split()
        tokens_kn = tokens_kn[0: min(len(tokens_kn), max_seq_length)]
    else: 
        tokens_kn = []

    if example.goal_text:
        tokens_goal = example.goal_text
        tokens_goal = tokens_goal.split()
        tokens_goal = tokens_goal[0: min(len(tokens_goal), max_seq_length)]
    else: 
        tokens_goal = []

    tokens_response = tokens_response[0: min(50, len(tokens_response))]
    if len(tokens_context) > max_seq_length - len(tokens_response) - 3: 
        tokens_context = tokens_context[len(tokens_context) \
                + len(tokens_response) - max_seq_length + 3:]

    context_tokens = []
    segment_ids = []

    context_tokens.append("[CLS]")
    segment_ids.append(0)
    context_tokens.extend(tokens_context)
    segment_ids.extend([0] * len(tokens_context))
    context_tokens.append("[SEP]")
    segment_ids.append(0)

    context_tokens.extend(tokens_response)
    segment_ids.extend([1] * len(tokens_response))
    context_tokens.append("[SEP]")
    segment_ids.append(1)

    context_ids = convert_tokens_to_ids(context_tokens)
    context_pos_ids = list(range(len(context_ids)))
    label_ids = label_map[example.label_text]
    if tokens_kn: 
        kn_ids = convert_tokens_to_ids(tokens_kn)
    else: 
        kn_ids = []

    if tokens_goal: 
        goal_ids = convert_tokens_to_ids(tokens_goal)
    else: 
        goal_ids = []

    feature = InputFeatures(
        context_ids=context_ids,
        context_pos_ids=context_pos_ids,
        segment_ids=segment_ids,
        kn_ids = kn_ids,
        goal_ids = goal_ids,
        label_ids=label_ids)
    #if ex_index < 5:
    #    print("*** Example ***")
    #    print("guid: %s" % (example.guid))
    #    print("context tokens: %s" % " ".join(context_tokens))
    #    print("context_ids: %s" % " ".join([str(x) for x in context_ids]))
    #    print("context_pos_ids: %s" % " ".join([str(x) for x in context_pos_ids]))
    #    print("segment_ids: %s" % " ".join([str(x) for x in segment_ids]))
    #    print("kn_ids: %s" % " ".join([str(x) for x in kn_ids]))
    #    print("label: %s (id = %d)" % (example.label_text, label_ids))
    return feature


#!/usr/bin/env python # -*- coding: utf-8 -*- ###################################################################### #   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. ###################################################################### File: data_provider.py load vocabulary dict generate self attention mask, [shape: batch_size *  max_len * max_len] pad_bath_kn_data Pad the instances to the max sequence length in batch, and generate the     corresponding position data and attention bias. process text to model inputs Base class for data converters for sequence classification data sets. Gets a collection of `InputExample`s for the train set. Gets a collection of `InputExample`s for the dev set. Gets a collection of `InputExample`s for prediction. Gets the list of labels for this data set. Converts a single `InputExample` into a single `InputFeatures`. generate_batch_data Reads a tab separated value file. Get number of examples for train, dev or test. Gets progress for training phase. Generate data for train, dev or test. instance_reader batch_reader wrapper A single training/test example input features datas Processor for the Match data set (GLUE version). See base class. See base class. See base class. See base class. Creates examples for the training and dev sets. convert input ids Converts a single `InputExample` into a single `InputFeatures`. #if ex_index < 5: #    print("*** Example ***") #    print("guid: %s" % (example.guid)) #    print("context tokens: %s" % " ".join(context_tokens)) #    print("context_ids: %s" % " ".join([str(x) for x in context_ids])) #    print("context_pos_ids: %s" % " ".join([str(x) for x in context_pos_ids])) #    print("segment_ids: %s" % " ".join([str(x) for x in segment_ids])) #    print("kn_ids: %s" % " ".join([str(x) for x in kn_ids])) #    print("label: %s (id = %d)" % (example.label_text, label_ids))