Fixed RAG sharding to include remainder of rows (facebookresearch#3836)

* Fixed RAG sharding to include remainder of rows * Less than not greater than
metacoding · Jul 22, 2021 · 53b8a8c · 53b8a8c
1 parent eae6d9c
commit 53b8a8c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/parlai/agents/rag/scripts/generate_dense_embeddings.py b/parlai/agents/rag/scripts/generate_dense_embeddings.py
@@ -182,6 +182,9 @@ def load_passages(self) -> List[Tuple[str, str, str]]:
         )
         shard_id, num_shards = self.opt['shard_id'], self.opt['num_shards']
         shard_size = int(len(rows) / num_shards)
+        if shard_id < len(rows) % num_shards:
+            # don't forget the remainder!
+            shard_size += 1
         start_idx = shard_id * shard_size
         end_idx = start_idx + shard_size
         logging.info(