Fix data yield

fe4e183f · Yibing Liu · 68e04f38 · fe4e183f · fe4e183f
隐藏空白更改
内联并排

Showing with 6 addition and 4 deletion

BERT/reader/cls.py BERT/reader/cls.py +3 -2

BERT/reader/squad.py BERT/reader/squad.py +3 -2

未找到文件。
--- a/BERT/reader/cls.py
+++ b/BERT/reader/cls.py
@@ -196,10 +196,11 @@ class DataProcessor(object):
                    return_num_token=False)
                if len(all_dev_batches) < dev_count:
                    all_dev_batches.append(batch_data)
-                else:
+
+                if len(all_dev_batches) == dev_count:
                    for batch in all_dev_batches:
                        yield batch
-                    all_dev_batches = [batch_data]
+                    all_dev_batches = []

        return wrapper


--- a/BERT/reader/squad.py
+++ b/BERT/reader/squad.py
@@ -566,10 +566,11 @@ class DataProcessor(object):
                        return_num_token=False)
                    if len(all_dev_batches) < dev_count:
                        all_dev_batches.append(batch_data)
-                    else:
+
+                    if len(all_dev_batches) == dev_count:
                        for batch in all_dev_batches:
                            yield batch
-                        all_dev_batches = [batch_data]
+                        all_dev_batches = []

        return wrapper