Enable Sem-dedup

[VibhuJawa]

Description

This PR builds on top #118 and adds the following features on top of it:

1. End to End Bash script
2. Improved Readme
3. Efficient and cleaned up compute_embeddings.py, clustering.py
4. Dask Accelerated sort_clusters.py and semdedup.py

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]

Thanks a ton @VibhuJawa, just a couple of nits and things I think you missed the first time around.

[faywang123]

I have tested the most recent PR (using 10 data files with 12 clusters). The result is consistent with our original result. Thanks, Vibhu! This is the command:

python semdedup_example.py --input-data-dir /ads_ds3/data/SemDeDup_BenchMark/datasets/c4/realnewslike/modified --config-file configs_cf.yml

The content of configs_cf.yml:

cache_dir: "/ads_ds3/data/SemDeDup_BenchMark"
num_files: 10
id_col_name: 'id' 
id_col_type: 'int' 
input_column: 'text'
input_file_type: 'json'
embeddings_save_loc: "/ads_ds3/data/SemDeDup_BenchMark/embeddings_fbopt_c4_10_pr130"
embedding_model_name_or_path: 'facebook/opt-125m' 
embedding_batch_size: 32
embedding_max_mem_gb: 10

clustering_save_loc: "/ads_ds3/data/SemDeDup_BenchMark/results_fbopt_c4_10_pr130"
n_clusters: 12 # -- number of clusters
seed: 1234
max_iter: 100
# Kmeans can only be done with L2 using cuML.
Kmeans_with_cos_dist: False # Only False allowed

# -- which example to keep from each group of duplicates
which_to_keep: "hard"
# largest cluster size the memory is large enough to process. If the
# cluster size is larger than it, we will devide the cluster into small
# clusters and process each one separately.
largest_cluster_size_to_process: 100000
sim_metric: "cosine" # Only cosine is allowed.
eps_thresholds:  0.001 0.01
eps_to_extract: 0.001

[VibhuJawa]

Thanks so much for this @faywang123 . Appreciate all the help.

@ayushdg , Can we use @faywang123 test above and put in our testing.

[ryantwolf]

Two nits then we're set.

[faywang123]

After a final review and test, the PR looks good to me. Thanks, @VibhuJawa for all the hard work!

[VibhuJawa]

@ryantwolf , Addressed the nits, let me know

[ryantwolf]

Incredible work, so excited to have this be a part of NeMo Curator

[ayushdg]

Not blocking but a couple of other suggestions:

1. Adding an optional SemDedup import to the top level modules/__init__.py file for gpu only environments. Allowing users to do something like from nemo_curator import SemDedup
2. Adding semantic deduplication in the list of features both in the README.md as well as a page in docs/user-guide

[VibhuJawa]

Not blocking but a couple of other suggestions:

1. Adding an optional SemDedup import to the top level modules/__init__.py file for gpu only environments. Allowing users to do something like from nemo_curator import SemDedup

Done .

2. Adding semantic deduplication in the list of features both in the README.md as well as a page in docs/user-guide

Added readme. docs/user-guide will be a followup.