Removal logic for fuzzy / exact (no class abstraction)

[praateekmahajan]

Description

See #529 for much more detailed analysis on the decision and followups

Removal Description

1. We implement ~left-anti join using a broadcast merge. This allows us to scale even when right is greater than device memory per node.
2. In our exact and fuzzy deduplication, we perform a shuffle at write time as the last step in identification. When reading for removal, we read FPP=1 so that we are reading a group at a time and not incurring "another" shuffle. By "another" shuffle I mean, that in a broadcast join, right is also reshuffled based on hash(id_col) so forcing a shuffle again will result in a double shuffle at read time.
3. We observe that the performance varies as num partitions (or partition sizes vary).
   • The fastest being where right is exactly one partition, in which case we broadcast all of right to all the workers, and perform a merge (resulting in least transfer).
4. One downside of this approach is when right is bigger than system memory per node, then the merge operation will have to spill, and that won't be possible, hence resulting in a stuck state where dask pauses workers once it reaches 80% capacity in which case user will have to implement a different merge logic.

Class Description

1. Adds a remove method to Fuzzy/ExactDeduplicator that calls nemo_curator.modules.removal.remove_duplicates
2. Moves logic from call to identify_duplicates.
3. Adds another paratmer to class construction for perform_removal which by default is False, to retain old behavior. But when True call removes the duplicates as well.

Usage

# Add snippet demonstrating usage

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[ryantwolf]

Looks good overall, one final change though can you modify nemo_curator/modules/__init__.py to include nemo_curator.modules.removal.remove_duplicates in the __all__? Model it after blend_datasets if you need a reference.

[sarahyurick]

Suggested change

	def identify(self, dataset: DocumentDataset) -> DocumentDataset:
	def _identify(self, dataset: DocumentDataset) -> DocumentDataset:

Nit, but maybe call them _identify and _remove if they are not intended to be accessed by the user directly.

[praateekmahajan]

Let's keep them exposed especially since remove won't work at scales where size of duplicate >> host memory, in which case the user will need to break down identify and remove

[sarahyurick]

Yes, that makes sense to me. What about calling it identify_duplicates?

[praateekmahajan]

Sounds good, initially I thought it's slightly verbose, but another argument in favor of identify_duplicates would be that in future we might want to expose identify_documents_to_keep in which the distinction might be necessary

cc @ayushdg / @ryantwolf / @VibhuJawa

[ryantwolf]

Looks good overall, just a few comments.

[sarahyurick]

Should also include the perform_removal option above?

[sarahyurick]

Yes, that makes sense to me. What about calling it identify_duplicates?

[sarahyurick]

Maybe call it "duplicates_removal" or something similar?

[VibhuJawa]

Thanks for pushing through this . This mostly looks good to me.

The only ask is not to Modify the __call__ header and behavior in this release. Everything else looks great.

[praateekmahajan]

@ayushdg we're doing this here

[ayushdg]

Left some initial comments. Would also be interested in reviewing the remove_duplicates utility once that's uploaded

[VibhuJawa]

LGTM (Assuming above reviews get through)

[sarahyurick]

Actually, I will cancel gpuCI since the last run looks good and only docs were updated.

[ayushdg]

Thanks lgtm!

Once we update the defaults to include removal as well, tutorials and CI scripts will need updates as well.