Commit 5622d832 authored by thomwolf's avatar thomwolf
Browse files

allowing to load small number of examples

parent a725db4f
......@@ -207,7 +207,7 @@ class InputFeatures(object):
self.end_position = end_position
def read_squad_examples(input_file, is_training):
def read_squad_examples(input_file, is_training, max_num=-1):
"""Read a SQuAD json file into a list of SquadExample."""
with tf.gfile.Open(input_file, "r") as reader:
input_data = json.load(reader)["data"]
......@@ -219,6 +219,8 @@ def read_squad_examples(input_file, is_training):
examples = []
for entry in input_data:
if max_num != -1 and len(examples) > max_num:
break
for paragraph in entry["paragraphs"]:
paragraph_text = paragraph["context"]
doc_tokens = []
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment