tqdm progress bar, smart num examples, random tmp files for output by default
@@ -19,8 +19,9 @@ from parlai.core.params import ParlaiParser
from parlai.agents.repeat_label.repeat_label import RepeatLabelAgent
from parlai.core.worlds import create_task
from parlai.core.utils import msg_to_str
-
import random
+import tempfile
+from tqdm import tqdm
def dump_data(opt):
@@ -28,10 +29,22 @@ def dump_data(opt):
agent = RepeatLabelAgent(opt)
world = create_task(opt, agent)
ignorefields = opt.get('ignore_fields', '')
+ if opt['outfile'] is None:
+ outfile = tempfile.mkstemp(
+ prefix='convert_data_to_parlai_output',
+ suffix='.txt')[1]
+ else:
+ outfile = opt['outfile']
+
+ if opt['num_examples'] == -1:
+ num_examples = world.num_examples()
+ else:
+ num_examples = opt['num_examples']
print('[ starting to convert.. ]')
- fw = open(opt['outfile'], 'w')
- for _ in range(opt['num_examples']):
+ print('[ saving output to {} ]'.format(outfile))
+ fw = open(outfile, 'w')
+ for _ in tqdm(range(num_examples)):
world.parley()
world.acts[0]['labels'] = world.acts[0].get(
'labels', world.acts[0].pop('eval_labels', None))
@@ -50,9 +63,15 @@ def main():
random.seed(42)
# Get command line arguments
parser = ParlaiParser()
- parser.add_argument('-n', '--num-examples', default=1000000000, type=int)
- parser.add_argument('-of', '--outfile', default='/tmp/dump', type=str)
- parser.add_argument('-if', '--ignore-fields', default='id', type=str)
+ parser.add_argument('-n', '--num-examples', default=-1, type=int,
+ help='Total number of exs to convert, -1 to convert \
+ all examples')
+ parser.add_argument('-of', '--outfile', default=None, type=str,
+ help='Output file where to save, by default will be \
+ created in /tmp')
+ parser.add_argument('-if', '--ignore-fields', default='id', type=str,
+ help='Ignore these fields from the message (returned\
+ with .act() )')
parser.set_defaults(datatype='train:stream')
opt = parser.parse_args()
dump_data(opt)
