Cleanup compute_embeddings_crossfit.py

Signed-off-by: Vibhu Jawa <[email protected]>

Author: VibhuJawa

nemo_curator/modules/semdedup/configs_cf.yml

@@ -1,29 +1,26 @@
-root: '/raid/fayw/semdedup'
-sample: 10
+root: '/ads_ds3/data/SemDeDup_BenchMark'
+sample: 12 # todo name to num_samples
 id_col:
-   name: 'id' #'adlr_id'
-   type:  'int' #'str'
-
+   name: 'adlr_id' # 'id' #'adlr_id'
+   type: 'str' # 'int' #'str'
 
 # Embeddings
 embeddings:
   input_column: 'text'
-  datapath: 'datasets/c4/realnewslike/modified/'
-  #datapath: 'datasets/prospector-lm/cleaned_exact_dedup_all_cc'
-  emb_parquet_path: "emb_fbopt125_c4_crossfit_10" #"emb_sentmpnet_c4_crossfit_5" # #
-  model_name: 'facebook/opt-125m' #'sentence-transformers/all-mpnet-base-v2' #  '/raid/fayw/semdedup/models/facebook-opt-125m/'
-
+  input_data_dir: '/datasets/prospector-lm/cleaned_exact_dedup_all_cc'
+  output_data_dir: '/raid/vjawa/prospector-lm/embeddings_crossfit_fb_c4_10'
+  path_or_name: 'facebook/opt-125m'
   emb_size: 768
-  batch_size: 32
-  num_workers: 8
-  pool_size: '10GB'
+  batch_size: 512
 
 clustering:
-  save_loc: "results_fbopt125_c4_crossfit_10" #"results_fbopt125_cc_crossfit_5"
-  num_clusters: 12 # -- number of clusters
+  save_loc: "results_fbopt_nc_11_crossfit" #"results_fbopt125_cc_crossfit_5"
+  #save_loc: "results_crossfit_fb_c4_10a"
+  num_clusters: 1000 # -- number of clusters
   seed: 1234
-  niter: 1000
-  Kmeans_with_cos_dist: True
+  niter: 100
+  # Kmeans can only be done with L2 using cuML.
+  Kmeans_with_cos_dist: False # Only False allowed
 
 
 semdedup:
@@ -33,9 +30,8 @@ semdedup:
   # 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'
-  keep_hard: True # True for hard examples
+  sim_metric: "cosine" # Only cosine is allowed.
 
 extract_dedup:
-  eps:  '0.01' # '0.01', '0.001', '0.0001', '1e-05', '1e-06'
-
+  use_eps_from_yml: True
+  eps:  0.01 0.001

setup.py

@@ -60,7 +60,8 @@
         "presidio-anonymizer==2.2.351",
         "usaddress==0.5.10",
         "nemo_toolkit[nlp]>=1.23.0",
-        "crossfit @ git+https://github.com/rapidsai/crossfit.git@1ee3de4",
+        "Cython",
+        "crossfit @ git+https://github.com/rapidsai/[email protected]",
         # justext installation breaks without lxml[html_clean]
         # due to this: https://github.com/miso-belica/jusText/issues/47
         "lxml[html_clean]",