_load_data as a generator (#4253)

facebookresearch · Dec 15, 2021 · d00f7d3 · d00f7d3
1 parent ce83b61
commit d00f7d3
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/parlai/tasks/wizard_of_internet/agents.py b/parlai/tasks/wizard_of_internet/agents.py
@@ -183,11 +183,9 @@ def add_cmdline_args(cls, parser: ParlaiParser, partial_opt=None) -> ParlaiParse
 
     def _load_data(self, datafile):
         logging.info(f'Loading data from {datafile} ...')
-        dialogs = []
         with jsonlines.open(datafile, 'r') as fin:
             for dialog_json in tqdm(fin):
-                dialogs.append(self._get_episode_examples(dialog_json))
-        return dialogs
+                yield self._get_episode_examples(dialog_json)
 
     def _get_episode_examples(self, dialog_json):
         data = get_single_val_from_dict(dialog_json)