data.py

import glob
import random
import struct
import copy
from tensorflow.core.example import example_pb2

PAD_TOKEN = '[PAD]' # This has a vocab id, which is used to pad the encoder input, decoder input and target sequence
UNKNOWN_TOKEN = '[UNK]' # This has a vocab id, which is used to represent out-of-vocabulary words
START_DECODING = '[START]' # This has a vocab id, which is used at the start of every decoder input sequence
STOP_DECODING = '[STOP]' # This has a vocab id, which is used at the end of untruncated target sequences

# Note: none of [PAD], [UNK], [START], [STOP] should appear in the vocab file.


class Vocab:
	"""Vocabulary class for mapping between words and ids (integers) """

	def __init__(self, vocab_file, max_size):
		"""Creates a vocab of up to max_size words, reading from the vocab_file. If max_size is 0, reads the entire vocab file.

		Args:
			vocab_file: path to the vocab file, which is assumed to contain "<word> <frequency>" on each line, sorted with most frequent word first. This code doesn't actually use the frequencies, though.
			max_size: integer. The maximum size of the resulting Vocabulary."""
		self._word_to_id = {}
		self._id_to_word = {}
		self._count = 0  # keeps track of total number of words in the Vocab

		# [PAD],[UNK], [START] and [STOP] get the ids 0,1,2,3.
		for w in [PAD_TOKEN, UNKNOWN_TOKEN, START_DECODING, STOP_DECODING]:
			self._word_to_id[w] = self._count
			self._id_to_word[self._count] = w
			self._count += 1

		# Read the vocab file and add words up to max_size
		with open(vocab_file, 'r', encoding="utf-8") as vocab_f:
			for line in vocab_f:
				pieces = line.split(" ")
				if len(pieces) != 2:
					print('Warning: incorrectly formatted line in vocabulary file: %s\n' % line)
					continue
				w = pieces[0]
				if w in [UNKNOWN_TOKEN, PAD_TOKEN, START_DECODING, STOP_DECODING]:
					raise Exception('[UNK], [PAD], [START] and [STOP] shouldn\'t be in the vocab file, but %s is' % w)
				if w in self._word_to_id:
					raise Exception('Duplicated word in vocabulary file: %s' % w)
				self._word_to_id[w] = self._count
				self._id_to_word[self._count] = w
				self._count += 1
				if max_size != 0 and self._count >= max_size:
					print("max_size of vocab was specified as %i; we now have %i words. Stopping reading." % (max_size, self._count))
					break

		print("Finished constructing vocabulary of %i total words. Last word added: %s" % (self._count, self._id_to_word[self._count-1]))

	def word2id(self, word):
		"""
		Returns the id (integer) of a word (string). Returns [UNK] id if word is OOV."""
		if word not in self._word_to_id:
			return self._word_to_id[UNKNOWN_TOKEN]
		return self._word_to_id[word]

	def id2word(self, word_id):
		"""
		Returns the word (string) corresponding to an id (integer)."""
		if word_id not in self._id_to_word:
			raise ValueError('Id not found in vocab: %d' % word_id)
		return self._id_to_word[word_id]

	def size(self):
		"""Returns the total size of the vocabulary"""
		return self._count


def example_generator(data_path, single_pass):
	"""Generates tf.Examples from data files.

		Binary data format: <length><blob>. <length> represents the byte size of <blob>.
		<blob> is serialized  tf.Example proto .
		The tf.Example contains the tokenized article text and summary.

	Args:
		data_path:
			Path to tf.Example data files. Can include wildcards( *), e.g. if you have several training data chunk files train_001.bin, train_002.bin, etc, then pass data_path=train_* to access them all.
		single_pass:
			Boolean. If True, go through the dataset exactly once, 只
			generating examples in the order they appear, then return. Otherwise, generate random examples indefinitely

	Yields:
		Deserialized tf.Example .
	"""
	epoch = 0
	while True:
		filelist = glob.glob(data_path)
		assert filelist, ('Error: Empty filelist at %s' % data_path)
		if single_pass:
			filelist = sorted(filelist)

		else:
			random.shuffle(filelist)
		for f in filelist:
			reader = open(f, 'rb')
			while True:
				len_bytes = reader.read(8)
				if not len_bytes:
					break # finished reading this file
				str_len = struct.unpack('q', len_bytes)[0]
				example_str = struct.unpack('%ds' % str_len, reader.read(str_len))[0]
				yield example_pb2.Example.FromString(example_str)
		if single_pass:
			print("example_generator completed reading all datafiles. No more data.")
			break

		epoch += 1


def background2ids(background_token, vocab):
	"""Map the article words to their ids. Also return a list of OOVs in the article.
	Args:
		background_token: list of words (strings)
		vocab: Vocabulary object

	Returns:
		ids:
			A list of word ids (integers); OOVs are represented by their temporary article OOV number. If the vocabulary size is 50k and the article has 3 OOVs, then these temporary OOV numbers will be 50000, 50001, 50002.
		oovs:
			A list of the OOV words in the article (strings), in the order corresponding to their temporary article OOV numbers. """
	ids = []
	oovs = []
	unk_id = vocab.word2id(UNKNOWN_TOKEN)
	for w in background_token:
		i = vocab.word2id(w)
		if i == unk_id:
			if w not in oovs:
				oovs.append(w)
			oov_num = oovs.index(w) # This is 0 for the first article OOV, 1 for the second article OOV...
			ids.append(vocab.size() + oov_num) # This is e.g. 50000 for the first article OOV, 50001 for the second...
		else:
			ids.append(i)
	return ids, oovs


def context2ids(context_token, vocab):
	"""Map the context words to their ids. Also return a list of OOVs in the context
	Args:
		context_token: list of words (strings)
		vocab: Vocabulary object

	Returns:
		ids:
			A list of word ids (integers); OOVs are represented by their temporary query OOV number. If the vocabulary size is 50k and the article has 3 OOVs, then these temporary OOV numbers will be 50000, 50001, 50002.
		oovs:
			A list of the OOV words in the article (strings), in the order corresponding to their temporary article OOV numbers."""
	ids = []
	oovs = []
	unk_id = vocab.word2id(UNKNOWN_TOKEN)
	for w in context_token:
		i = vocab.word2id(w)
		if i == unk_id: # If w is OOV
			if w not in oovs: # Add to list of OOVs
				oovs.append(w)
			oov_num = oovs.index(w) # This is 0 for the first article OOV, 1 for the second article OOV...
			ids.append(vocab.size() + oov_num) # This is e.g. 50000 for the first article OOV, 50001 for the second...
		else:
			ids.append(i)
	return ids, oovs


def response2ids(response_token , vocab, background_oovs):
	"""Map the abstract words to their ids. In-article OOVs are mapped to their temporary OOV numbers.

	Args:
		response_token: list of words (strings)
		vocab: Vocabulary object
		background_oovs: list of in-article OOV words (strings), in the order corresponding to their temporary article OOV numbers

	Returns:
		ids: List of ids (integers). In-article OOV words are mapped to their temporary OOV numbers. Out-of-article OOV words are mapped to the UNK token id."""
	ids = []
	unk_id = vocab.word2id(UNKNOWN_TOKEN)
	for w in response_token:
		i = vocab.word2id(w)
		if i == unk_id: # If w is an OOV word
			if w in background_oovs: # If w is an in-article OOV
				vocab_idx = vocab.size() + background_oovs.index(w) # Map to its temporary article OOV number
				ids.append(vocab_idx)
			else:  # If w is an out-of-article OOV
				ids.append(unk_id)
		else:
			ids.append(i)
	return ids


def outputids2words(id_list, vocab, background_oovs, backgrounds_token): # response in our case
	"""Maps output ids to words, including mapping in-article OOVs from their temporary ids to the original OOV string.
	Args:
		id_list: list of ids (integers)
		vocab: Vocabulary object
		background_oovs: list of OOV words (strings) in the order corresponding to their temporary article OOV ids
	Returns:
		words: list of words (strings)
	"""
	words = []
	highlights = []
	spans =[]
	for i in id_list:
		if isinstance(i, list):
			words = words + backgrounds_token[i[0]:(i[1]+1)]
			spans = spans + backgrounds_token[i[0]:(i[1]+1)]
			highlights = highlights + ["<start<"] + backgrounds_token[i[0]:(i[1]+1)] + [">end>"]
		else:
			try:
				w = vocab.id2word(i)
			except ValueError as e:
				assert background_oovs is not None, "Error: model produced a word ID that isn't in the vocabulary. This should not happen in baseline (no pointer-generator) mode"
				background_oov_idx = i - vocab.size()
				try:
					w = background_oovs[background_oov_idx]
				except IndexError as e:  # i doesn't correspond to an article oov
					raise ValueError('Error: model produced word ID %i which corresponds to article OOV %i but this example only has %i article OOVs' % (i, background_oov_idx, len(background_oovs)))
			words.append(w)
			highlights.append(w)
	return words, highlights,spans


def show_background_span(background_token, b_start, b_end):
	modify_background_token = copy.copy(background_token)
	modify_background_token.insert(b_start, "<start<")
	modify_background_token.insert(b_end+2, ">end>")
	out_str = ' '.join(modify_background_token)
	return out_str


def show_response_span(response_token, r_start, r_end):
	modify_response_token = copy.copy(response_token)
	modify_response_token.insert(r_start, "<start<")
	modify_response_token.insert(r_end+2, ">end>")
	out_str = ' '.join(modify_response_token)
	return out_str


def show_background_oovs(background_text, vocab):
	"""Returns the article string, highlighting the OOVs by placing __underscores__ around them"""
	unk_token = vocab.word2id(UNKNOWN_TOKEN)
	words = background_text.split(' ')
	words = [("__%s__" % w) if vocab.word2id(w)==unk_token else w for w in words]
	out_str = ' '.join(words)
	return out_str


def show_context_oovs(context_text, vocab):
	"""Returns the query string, highlighting the OOVs by placing __underscores__ around them"""
	unk_token = vocab.word2id(UNKNOWN_TOKEN)
	words = context_text.split(' ')
	words = [("__%s__" % w) if vocab.word2id(w)==unk_token else w for w in words]
	out_str = ' '.join(words)
	return out_str


def show_response_oovs(response_text, vocab, background_oovs):
	"""Returns the abstract string, highlighting the article OOVs with __underscores__.

	If a list of article_oovs is provided, non-article OOVs are differentiated like !!__this__!!.

	Args:
		response_text: string
		vocab: Vocabulary object
		background_oovs: list of words (strings)
	"""
	unk_token = vocab.word2id(UNKNOWN_TOKEN)
	words = response_text.split(' ')
	new_words = []
	for w in words:
		if vocab.word2id(w) == unk_token: # w is oov
			if background_oovs is None: # baseline mode
				new_words.append("__%s__" % w)
			else: # pointer-generator mode
				if w in background_oovs:
					new_words.append("__%s__" % w)
				else:
					new_words.append("!!__%s__!!" % w)
		else: # w is in-vocab word
			new_words.append(w)
	out_str = ' '.join(new_words)
	return out_str