fix a bug function _indexize in fastNLP/io/pipe/utils.py

fastNLP/io/pipe/utils.py

@@ -105,18 +105,20 @@ def _indexize(data_bundle, input_field_names=Const.INPUT, target_field_names=Con
         target_field_names = [target_field_names]
     for input_field_name in input_field_names:
         src_vocab = Vocabulary()
-        src_vocab.from_dataset(data_bundle.datasets['train'], field_name=input_field_name,
-                               no_create_entry_dataset=[dataset for name, dataset in data_bundle.datasets.items() if
-                                                        name != 'train'])
+        src_vocab.from_dataset(*[ds for name, ds in data_bundle.iter_datasets() if 'train' in name],
+                               field_name=input_field_name,
+                               no_create_entry_dataset=[ds for name, ds in data_bundle.iter_datasets()
+                                                        if ('train' not in name) and (ds.has_field(input_field_name))]
+                               )
         src_vocab.index_dataset(*data_bundle.datasets.values(), field_name=input_field_name)
         data_bundle.set_vocab(src_vocab, input_field_name)
     
     for target_field_name in target_field_names:
         tgt_vocab = Vocabulary(unknown=None, padding=None)
         tgt_vocab.from_dataset(*[ds for name, ds in data_bundle.iter_datasets() if 'train' in name],
-                               field_name=Const.TARGET,
+                               field_name=target_field_name,
                                no_create_entry_dataset=[ds for name, ds in data_bundle.iter_datasets()
-                                                        if ('train' not in name) and (ds.has_field(Const.TARGET))]
+                                                        if ('train' not in name) and (ds.has_field(target_field_name))]
                                )
         if len(tgt_vocab._no_create_word) > 0:
             warn_msg = f"There are {len(tgt_vocab._no_create_word)} target labels" \