Better handling of unspecified max_tokens and max_sentences (fixes #1427)

Summary: Pull Request resolved: https://github.com/fairinternal/fairseq-py/pull/928

Differential Revision: D18691525

Pulled By: myleott

fbshipit-source-id: e787c17434d4cb0c4621e9858e0ebec4f9951630

fairseq/data/data_utils.py

@@ -221,8 +221,8 @@ def batch_by_size(
             'or `python setup.py build_ext --inplace`'
         )
 
-    max_tokens = max_tokens if max_tokens is not None else sys.maxsize
-    max_sentences = max_sentences if max_sentences is not None else sys.maxsize
+    max_tokens = max_tokens if max_tokens is not None else -1
+    max_sentences = max_sentences if max_sentences is not None else -1
     bsz_mult = required_batch_size_multiple
 
     if isinstance(indices, types.GeneratorType):

fairseq/data/data_utils_fast.pyx

@@ -16,9 +16,9 @@ ctypedef np.int64_t DTYPE_t
 cdef _is_batch_full(list batch, long num_tokens, long max_tokens, long max_sentences):
     if len(batch) == 0:
         return 0
-    if len(batch) == max_sentences:
+    if max_sentences > 0 and len(batch) == max_sentences:
         return 1
-    if num_tokens > max_tokens:
+    if max_tokens > 0 and num_tokens > max_tokens:
         return 1
     return 0
 
@@ -47,7 +47,7 @@ cpdef list batch_by_size_fast(
         sample_lens.append(num_tokens)
         sample_len = max(sample_len, num_tokens)
 
-        assert sample_len <= max_tokens, (
+        assert max_tokens <= 0 or sample_len <= max_tokens, (
             "sentence at index {} of size {} exceeds max_tokens "
             "limit of {}!".format(idx, sample_len, max_tokens)
         )