import torchtext.data as data
from data.torchtext.custom_datasets import Dataset

class CustomSequenceTaggingDataSet(Dataset):
	"""Defines a dataset for sequence tagging. Examples in this dataset
	contain paired lists -- paired list of words and tags.

	For example, in the case of part-of-speech tagging, an example is of the
	form
	[I, love, PyTorch, .] paired with [PRON, VERB, PROPN, PUNCT]

	See torchtext/test/sequence_tagging.py on how to use this class.
	"""

	@staticmethod
	def sort_key(example):
		for attr in dir(example):
			if not callable(getattr(example, attr)) and \
					not attr.startswith("__"):
				return len(getattr(example, attr))
		return 0

	def __init__(self, path, fields, separator="\t", **kwargs):
		examples = []
		columns = []

		with open(path) as input_file:
			for line in input_file:
				if line.startswith("-DOCSTART-"):
					continue
				
				line = line.strip()
				if line == "":
					if columns:
						# copy first column as a complete sentence that is not tokenized
						#sentence = columns[0].copy()
						#columns.append(sentence)
						examples.append(data.Example.fromlist(columns, fields))
					columns = []
				else:
					for i, column in enumerate(line.split(separator)):
						if len(columns) < i + 1:
							columns.append([])
						column = column.replace('I-', '')
						column = column.replace('B-', '')

						columns[i].append(column)

			if columns:
				examples.append(data.Example.fromlist(columns, fields))
		super(CustomSequenceTaggingDataSet, self).__init__(examples, fields,
													 **kwargs)
