Explore the tone-check data_generation_templates notebook created by @achou and incorporate its logic into the tone-check training data generation job in ml-pipelines. Create a tone-check data generation DAG that uses the job logic artifact and runs in the airflow-ml instance.
Description
Description
Details
Details
Related Changes in GitLab:
| Title | Reference | Author | Source Branch | Dest Branch | |
|---|---|---|---|---|---|
| ml: add tone-check data generation DAG | repos/data-engineering/airflow-dags!1699 | kevinbazira | tone_check_data_generation_dag | main | |
| tone-check: add data generation pipeline job logic | repos/machine-learning/ml-pipelines!52 | kevinbazira | tone_check_data_generation | main |
Event Timeline
Comment Actions
I ended up removing the dev-limits as the small sample size results in no rows making it through the end of the pipeline as shown below:
Job running with dev-limit at 100000
$ yarn logs -appOwner kevinbazira -applicationId application_1757517622464_118136 | grep " - INFO - " Picked up JAVA_TOOL_OPTIONS: -Dfile.encoding=UTF-8 25/09/16 02:09:12 INFO zlib.ZlibFactory: Successfully loaded & initialized native-zlib library 25/09/16 02:09:12 INFO compress.CodecPool: Got brand-new decompressor [.deflate] 2025-09-15 17:14:45,022 - INFO - Starting tone_check_data_generation job. 2025-09-15 17:14:45,023 - INFO - Parameters: output-path=/user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/tone_check_training_data.parquet, wiki-db=enwiki, snapshot=2025-04 2025-09-15 17:14:45,023 - INFO - Development mode: Limiting data to first 100000 revisions. 2025-09-15 17:14:52,433 - INFO - Spark Session created successfully. 2025-09-15 17:14:52,433 - INFO - Step 1: Reading mediawiki_content_history_v1 and matching templates. 2025-09-15 17:14:54,745 - INFO - Development mode: Limited content_history to 100000 rows. 2025-09-15 17:33:52,549 - INFO - Template matching complete. Number of revisions with templates: 757 2025-09-15 17:53:53,779 - INFO - Identified 13 unique pages with relevant templates. 2025-09-15 17:53:53,780 - INFO - Joining template data with mediawiki_history. 2025-09-15 17:53:54,210 - INFO - Development mode: Limited mediawiki_history to 100000 rows. 2025-09-15 18:15:33,187 - INFO - Joined with mediawiki_history. Number of rows: 1 2025-09-15 18:56:10,753 - INFO - Joined with template metadata. Number of rows: 1 2025-09-15 18:56:10,760 - INFO - Step 2: Extracting positive and negative revision pairs. 2025-09-15 19:26:16,326 - INFO - Removed reverted revisions. Number of rows: 0 2025-09-15 20:11:30,687 - INFO - Filtered pages with at least one template. Number of pages: 0 2025-09-15 20:53:54,392 - INFO - Filtered revisions for pages with templates. Number of rows: 0 2025-09-15 21:31:57,048 - INFO - Filtered revisions with template change. Number of rows: 0 2025-09-15 22:05:02,487 - INFO - Final positive/negative revision pairs. Number of rows: 0 2025-09-15 22:05:02,491 - INFO - Step 3: Joining with revision content to add wikitext. 2025-09-15 22:07:11,421 - INFO - Loaded content_history for join. Number of rows: 848036146 2025-09-15 22:44:35,958 - INFO - Joined with revision content. Number of rows: 0 2025-09-15 22:44:36,063 - INFO - Final DataFrame schema created. 2025-09-15 22:44:36,074 - INFO - Starting final write to: /user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/tone_check_training_data.parquet 2025-09-15 23:32:44,091 - INFO - Write operation complete. 2025-09-16 00:08:37,645 - INFO - Wrote 0 rows to /user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/tone_check_training_data.parquet 2025-09-16 00:08:41,295 - INFO - Job finished successfully.
Job running without dev-limit
$ yarn logs -appOwner kevinbazira -applicationId application_1757517622464_128337 | grep " - INFO - " Picked up JAVA_TOOL_OPTIONS: -Dfile.encoding=UTF-8 25/09/16 14:12:37 INFO zlib.ZlibFactory: Successfully loaded & initialized native-zlib library 25/09/16 14:12:37 INFO compress.CodecPool: Got brand-new decompressor [.deflate] 2025-09-16 04:10:14,664 - INFO - Starting tone_check_data_generation job. 2025-09-16 04:10:14,664 - INFO - Parameters: output-path=/user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/tone_check_training_data.parquet, wiki-db=enwiki, snapshot=2025-04 2025-09-16 04:10:33,251 - INFO - Spark Session created successfully. 2025-09-16 04:10:33,251 - INFO - Step 1: Reading mediawiki_content_history_v1 and matching templates. 2025-09-16 04:27:20,061 - INFO - Template matching complete. Number of revisions with templates: 3784610 2025-09-16 04:45:28,521 - INFO - Identified 103300 unique pages with relevant templates. 2025-09-16 04:45:28,521 - INFO - Joining template data with mediawiki_history. 2025-09-16 05:05:15,594 - INFO - Joined with mediawiki_history. Number of rows: 54126103 2025-09-16 05:42:25,828 - INFO - Joined with template metadata. Number of rows: 54126103 2025-09-16 05:42:25,834 - INFO - Step 2: Extracting positive and negative revision pairs. 2025-09-16 06:35:40,079 - INFO - Removed reverted revisions. Number of rows: 36534608 2025-09-16 07:10:39,391 - INFO - Filtered pages with at least one template. Number of pages: 96329 2025-09-16 08:00:57,497 - INFO - Filtered revisions for pages with templates. Number of rows: 32090267 2025-09-16 08:52:53,160 - INFO - Filtered revisions with template change. Number of rows: 223252 2025-09-16 09:49:22,879 - INFO - Final positive/negative revision pairs. Number of rows: 211954 2025-09-16 09:49:22,885 - INFO - Step 3: Joining with revision content to add wikitext. 2025-09-16 09:51:30,691 - INFO - Loaded content_history for join. Number of rows: 848199540 2025-09-16 10:49:17,641 - INFO - Joined with revision content. Number of rows: 211950 2025-09-16 10:49:17,846 - INFO - Final DataFrame schema created. 2025-09-16 10:49:17,855 - INFO - Starting final write to: /user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/tone_check_training_data.parquet 2025-09-16 12:54:16,501 - INFO - Write operation complete. 2025-09-16 13:49:02,001 - INFO - Wrote 211950 rows to /user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/tone_check_training_data.parquet 2025-09-16 13:49:03,738 - INFO - Job finished successfully.
Comment Actions
As shown in T404722#11185284, the job running without a dev-limit finally succeeded after >9.5hrs. Below are the results from the generated tone_check_training_data.parquet:
parquet_df.show(10) +-------+--------+--------------------+-----------+---------------+--------------------+-----+ |wiki_db| page_id| page_title|revision_id|revision_id_key| clean_text|label| +-------+--------+--------------------+-----------+---------------+--------------------+-----+ | enwiki|43692302| BOL_Network| 1157721672| 1157705901|BOL Network () is...| 0| | enwiki| 6639998| Qooxdoo| 185014561| 177069289|qooxdoo () is an ...| 0| | enwiki|33446745| A_Common_Man_(film)| 504973054| 504393803|A Common Man is a...| 0| | enwiki| 2542010| Chopard| 201322196| 225050798|Chopard is a luxu...| 1| | enwiki|28115576| Evgeny_Vinokurov| 963364130| 1056567551|Vinokurov, Evgeny...| 1| | enwiki|17224876| Romina_Arena| 728133288| 961924766|Romina Arena (bor...| 1| | enwiki| 3509454| Treponem_Pal| 1122401755| 945909857|Treponem Pal is a...| 0| | enwiki|32379966| AL_and_AL| 491850708| 494128288|Al Holmes (born J...| 1| | enwiki|58396295| Caroline_Currid| 1099452687| 861964850|Caroline Currid (...| 0| | enwiki| 179979|The_College_of_Ne...| 604523151| 604522937|The College of Ne...| 0| +-------+--------+--------------------+-----------+---------------+--------------------+-----+ only showing top 10 rows
The model training pod in prod (T398970#11167313) was built using synthetic_tone_check_training_data.csv, which has the following format:
csv_df.show(10) +--------------------+-----+ | input|label| +--------------------+-----+ |en[SEP]Peacock_De...| 0| |en[SEP]Peacock_De...| 1| |en[SEP]Peacock_De...| 0| |en[SEP]Peacock_De...| 1| |en[SEP]Peacock_De...| 0| |en[SEP]Peacock_De...| 1| |en[SEP]Peacock_De...| 0| |en[SEP]Peacock_De...| 1| |en[SEP]Peacock_De...| 0| |en[SEP]Peacock_De...| 1| +--------------------+-----+ only showing top 10 rows
As a next step, I will add functionality to transform tone_check_training_data.parquet to tone_check_training_data.csv with the format: lang[SEP]page_title[SEP]clean_text in the input column.
Comment Actions
@kevinbazira Alternatively we could stick with parquet which is more efficient in terms of storage, so we could adapt the data loading part in the training phase. wdyt?
Comment Actions
Hi! I want to add more information regarding the data used for training. I was checking the original dataset for the peacock model. The dataset are split into 3 separate files: peacock_train.csv.gz along with peacock_val.csv.gz for training, and peacock_test.csv.gz for evaluation.
Key stats for 'peacock_train.csv.gz' and 'peacock_val.csv.gz':
- 7 languages (with varying sample sizes per language)
- Balanced sample distribution across labels
Here is a code snippet of how to calculate the stats (taken from original training notebook)
I think Aitolkyn sampled from the complete generated parquet file and then saved the train/val/test splits as separate csv.gz files. We could do the same here, because the parquet file likely contains a much larger dataset than we actually need for training (like we only need ~9k for enwiki).
Comment Actions
@isarantopoulos we initially planned to convert the parquet to csv because that's the format the model training logic expects. Since sticking with parquet would normally be the ideal approach, we are going to keep the data in parquet format. Here are the results from the model-ready tone_check_training_data.parquet:
model_ready_df.show(10) +--------------------+-----+ | input|label| +--------------------+-----+ |en[SEP]GameFront[...| 0| |en[SEP]Jamia_Mill...| 0| |en[SEP]Will_Duran...| 1| |en[SEP]MAKOplasty...| 0| |en[SEP]Michael_La...| 0| |en[SEP]Arcos_Dora...| 0| |en[SEP]What_Goes_...| 0| |en[SEP]Hakomi[SEP...| 1| |en[SEP]Susanne_Bo...| 0| |en[SEP]Unicode_Co...| 1| +--------------------+-----+ only showing top 10 rows
We will update the model training job logic to read from parquet instead of csv.
Comment Actions
Thank you for sharing the additional information, @achou. I have updated the job logic to use a reproducible random split (80/10/10) to create train_data.parquet, validation_data.parquet, and test_data.parquet from the model-ready data. Below are the results:
model_ready_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/model_ready_tone_check_training_data.parquet")
model_ready_df.count()
211950
model_ready_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]GameFront[...| 0|
|en[SEP]Jamia_Mill...| 0|
|en[SEP]Will_Duran...| 1|
|en[SEP]MAKOplasty...| 0|
|en[SEP]Michael_La...| 0|
|en[SEP]Arcos_Dora...| 0|
|en[SEP]What_Goes_...| 0|
|en[SEP]Hakomi[SEP...| 1|
|en[SEP]Susanne_Bo...| 0|
|en[SEP]Unicode_Co...| 1|
+--------------------+-----+
only showing top 10 rows
model_ready_df.groupBy("label").count().show()
+-----+------+
|label| count|
+-----+------+
| 1|105977|
| 0|105973|
+-----+------+
train_data_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/train_data.parquet")
train_data_df.count()
169680
train_data_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]The_Biodes...| 0|
|en[SEP]Kathy_Davi...| 0|
|en[SEP]Charles_Gh...| 1|
|en[SEP]Islam_and_...| 1|
|en[SEP]King_Mandh...| 1|
|en[SEP]Gawr_Gura[...| 1|
|en[SEP]Factory_Th...| 0|
|en[SEP]Dialog_Axi...| 0|
|en[SEP]William_Sc...| 1|
|en[SEP]Saskatchew...| 1|
+--------------------+-----+
only showing top 10 rows
train_data_df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
| 1|84974|
| 0|84706|
+-----+-----+
validation_data_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/validation_data.parquet")
validation_data_df.count()
21057
validation_data_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]Midtown_Co...| 1|
|en[SEP]Bobby_Kert...| 0|
|en[SEP]Loyola_Mar...| 1|
|en[SEP]Tigerair_A...| 0|
|en[SEP]TD_Synnex[...| 1|
|en[SEP]Harvey_Mud...| 0|
|en[SEP]Transcende...| 0|
|en[SEP]Research_i...| 0|
|en[SEP]Illumina,_...| 1|
|en[SEP]Aviad_Klei...| 0|
+--------------------+-----+
only showing top 10 rows
validation_data_df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
| 1|10423|
| 0|10634|
+-----+-----+
test_data_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/data/etl_output/test_data.parquet")
test_data_df.count()
21213
test_data_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]Threat_Sta...| 1|
|en[SEP]Allan_Loeb...| 0|
|en[SEP]Oracle_App...| 0|
|en[SEP]Upper_Colu...| 0|
|en[SEP]Disney's_A...| 0|
|en[SEP]Safra_Grou...| 0|
|en[SEP]NewsX[SEP]...| 1|
|en[SEP]Wolfgang_A...| 1|
|en[SEP]Zico_(foot...| 0|
|en[SEP]McGuffey_S...| 1|
+--------------------+-----+
only showing top 10 rows
test_data_df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
| 1|10580|
| 0|10633|
+-----+-----+Comment Actions
As mentioned in T404722#11188124, the tone-check training data generation job completed after running for >9.5hrs. At that time, it was using ~3.5TB of peak memory ((8G driver + 20G executor) * 128 executors).
Following this slack conversation, we agreed to reduce the memory resources to avoid overloading the cluster. I halved the resources and reran the pipeline with a peak memory usage of ~1.7TB ((8G driver + 20G executor) * 64 executors). This run took >20.5hrs and successfully wrote the full model-ready data to HDFS.
However, I noticed that the job did not finish splitting the model-ready data into training, validation, and test sets, as shown in the logs below:
$ yarn logs -appOwner kevinbazira -applicationId application_1757517622464_199779 | grep " - INFO - " Picked up JAVA_TOOL_OPTIONS: -Dfile.encoding=UTF-8 Unable to get ApplicationState. Attempting to fetch logs directly from the filesystem. 25/09/22 05:31:12 INFO zlib.ZlibFactory: Successfully loaded & initialized native-zlib library 25/09/22 05:31:12 INFO compress.CodecPool: Got brand-new decompressor [.deflate] 2025-09-19 04:23:16,590 - INFO - Starting tone_check_data_generation job. 2025-09-19 04:23:16,590 - INFO - Parameters: output-path=/user/kevinbazira/tone_check_data_generation_pipeline/20250919T042229, wiki-db=enwiki, snapshot=2025-04 2025-09-19 04:23:39,462 - INFO - Spark Session created successfully. 2025-09-19 04:23:39,462 - INFO - Step 1: Reading mediawiki_content_history_v1 and matching templates. 2025-09-19 05:01:11,508 - INFO - Template matching complete. Number of revisions with templates: 3785267 2025-09-19 05:43:07,983 - INFO - Identified 103303 unique pages with relevant templates. 2025-09-19 05:43:07,983 - INFO - Joining template data with mediawiki_history. 2025-09-19 06:28:06,033 - INFO - Joined with mediawiki_history. Number of rows: 54131396 2025-09-19 07:39:07,308 - INFO - Joined with template metadata. Number of rows: 54131396 2025-09-19 07:39:07,315 - INFO - Step 2: Extracting positive and negative revision pairs. 2025-09-19 08:58:46,607 - INFO - Removed reverted revisions. Number of rows: 36535387 2025-09-19 10:07:07,454 - INFO - Filtered pages with at least one template. Number of pages: 96324 2025-09-19 11:51:12,619 - INFO - Filtered revisions for pages with templates. Number of rows: 32089938 2025-09-19 13:49:08,616 - INFO - Filtered revisions with template change. Number of rows: 223238 2025-09-19 15:39:49,972 - INFO - Final positive/negative revision pairs. Number of rows: 211940 2025-09-19 15:39:49,974 - INFO - Step 3: Joining with revision content to add wikitext. 2025-09-19 15:46:50,540 - INFO - Loaded content_history for join. Number of rows: 848439527 2025-09-19 17:56:50,284 - INFO - Joined with revision content. Number of rows: 211936 2025-09-19 17:56:50,382 - INFO - Transforming structured data into format: input and label columns, with input as lang[SEP]page_title[SEP]clean_text. 2025-09-19 17:56:50,405 - INFO - model-ready DataFrame schema created (input/label format). 2025-09-19 17:56:50,416 - INFO - Writing full model-ready dataset to: /user/kevinbazira/tone_check_data_generation_pipeline/20250919T042229/full_model_ready_data.parquet 2025-09-20 01:22:05,187 - INFO - Full model-ready dataset write complete. Number of rows: 211936 2025-09-20 01:22:05,190 - INFO - Step 4: Splitting model-ready data into train, validation, and test sets (80/10/10 split). 2025-09-20 01:22:05,236 - INFO - Writing train set to: /user/kevinbazira/tone_check_data_generation_pipeline/20250919T042229/train_data.parquet
I am going to investigate why the job did not complete the splitting step and why the three separate files for training, validation, and testing were not fully written.
Comment Actions
As part of investigating why the DAG in T404722#11200120 failed to complete the data splitting step, I wanted to confirm whether reduced resource allocation was the cause of this issue. To test this, I isolated the splitting logic into a standalone script (P83449) and ran it manually against the full model-ready dataset. The results of this test were:
1.With default memory configurations, the split completed in <5 minutes:
spark3-submit \ --master yarn \ split_model_ready_data.py \ --input-path $HDFS_INPUT_FILE_PATH \ --output-path $HDFS_OUTPUT_FILE_PATH
2.With memory settings matching the DAG, the job completed in <2 minutes:
spark3-submit \ --master yarn \ --driver-cores 4 \ --driver-memory 8G \ --executor-cores 4 \ --executor-memory 20G \ --conf spark.dynamicAllocation.maxExecutors=64 \ --conf spark.executor.memoryOverhead=8G \ --conf spark.driver.memoryOverhead=8G \ --conf spark.driver.maxResultSize=16G \ --conf spark.sql.shuffle.partitions=16384 \ --conf spark.sql.autoBroadcastJoinThreshold=-1 \ --conf spark.sql.adaptive.enabled=true \ --conf spark.network.timeout=1200s \ split_model_ready_data.py \ --input-path $HDFS_INPUT_FILE_PATH \ --output-path $HDFS_OUTPUT_FILE_PATH
I also reviewed the failed DAG's YARN error logs for key error signatures:
There were no OOM issues in the logs, and the manual split completed quickly even with the same memory configuration as the DAG. This confirms that the problem is not related to memory resources, but is likely due to how the DAG process is managed or how the job is executed within Airflow.
Since we know that the model-ready data generation step runs without issues and the data splitting step completes successfully when isolated, as a next step, I am going to refactor the DAG into two separate, sequential tasks:
- generate_training_data: job responsible only for the data generation and producing the model-ready data parquet file.
- split_training_data: job that takes the model-ready data file as input and splits it into train, validation, and test sets.
Comment Actions
After separating the training data generation and data splitting steps into sequential DAG tasks as discussed in T404722#11201309, the tone-check data generation pipeline has now run end-to-end successfully:
1.generate_training_data: job completed in ~6.5hrs:
$ yarn logs -appOwner kevinbazira -applicationId application_1758550516505_0971 | grep " - INFO - " Picked up JAVA_TOOL_OPTIONS: -Dfile.encoding=UTF-8 Unable to get ApplicationState. Attempting to fetch logs directly from the filesystem. 25/09/23 05:18:59 INFO zlib.ZlibFactory: Successfully loaded & initialized native-zlib library 25/09/23 05:18:59 INFO compress.CodecPool: Got brand-new decompressor [.deflate] 2025-09-22 15:17:33,129 - INFO - Starting tone_check_data_generation job. 2025-09-22 15:17:33,130 - INFO - Parameters: output-path=/user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/full_model_ready_data.parquet, wiki-db=enwiki, snapshot=2025-04 2025-09-22 15:17:48,475 - INFO - Spark Session created successfully. 2025-09-22 15:17:48,475 - INFO - Step 1: Reading mediawiki_content_history_v1 and matching templates. 2025-09-22 15:17:50,896 - INFO - Template matching complete. 2025-09-22 15:17:50,907 - INFO - Identified unique pages with relevant templates. 2025-09-22 15:17:50,907 - INFO - Joining template data with mediawiki_history. 2025-09-22 15:17:51,207 - INFO - Joined with mediawiki_history. 2025-09-22 15:17:51,227 - INFO - Joined with template metadata. 2025-09-22 15:17:51,227 - INFO - Step 2: Extracting positive and negative revision pairs. 2025-09-22 15:17:51,253 - INFO - Removed reverted revisions. 2025-09-22 15:17:51,268 - INFO - Filtered pages with at least one template. 2025-09-22 15:17:51,313 - INFO - Filtered revisions for pages with templates. 2025-09-22 15:17:51,360 - INFO - Filtered revisions with template change. 2025-09-22 15:17:51,505 - INFO - Final positive/negative revision pairs. 2025-09-22 15:17:51,505 - INFO - Step 3: Joining with revision content to add wikitext. 2025-09-22 15:17:51,524 - INFO - Loaded content_history for join. 2025-09-22 15:17:51,534 - INFO - Joined with revision content. 2025-09-22 15:17:51,580 - INFO - Transforming structured data into format: input and label columns, with input as lang[SEP]page_title[SEP]clean_text. 2025-09-22 15:17:51,600 - INFO - model-ready DataFrame schema created (input/label format). 2025-09-22 15:17:51,610 - INFO - Writing full model-ready dataset to: /user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/full_model_ready_data.parquet 2025-09-22 22:03:24,274 - INFO - Full model-ready dataset write complete. Number of rows: 211926 2025-09-22 22:03:27,001 - INFO - Data generation job finished successfully.
2.split_training_data: job finished in ~1.5mins:
$ yarn logs -appOwner kevinbazira -applicationId application_1758550516505_7426 | grep " - INFO - " Picked up JAVA_TOOL_OPTIONS: -Dfile.encoding=UTF-8 Unable to get ApplicationState. Attempting to fetch logs directly from the filesystem. 25/09/23 05:23:58 INFO zlib.ZlibFactory: Successfully loaded & initialized native-zlib library 25/09/23 05:23:58 INFO compress.CodecPool: Got brand-new decompressor [.deflate] 2025-09-22 22:04:21,795 - INFO - Reading model-ready parquet data from: /user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/full_model_ready_data.parquet 2025-09-22 22:04:45,301 - INFO - Splitting model-ready data into train, validation, and test sets (80/10/10 split). 2025-09-22 22:04:45,344 - INFO - Writing train set to: /user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/train_data.parquet 2025-09-22 22:05:13,384 - INFO - Train set write complete. Number of rows: 169749 2025-09-22 22:05:13,384 - INFO - Writing validation set to: /user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/validation_data.parquet 2025-09-22 22:05:30,420 - INFO - Validation set write complete. Number of rows: 21114 2025-09-22 22:05:30,421 - INFO - Writing test set to: /user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/test_data.parquet 2025-09-22 22:05:47,015 - INFO - Test set write complete. Number of rows: 21023 2025-09-22 22:05:47,637 - INFO - Data split job finished successfully.
NOTE: Removed unnecessary DataFrame .count() operations that were previously used for logging/debugging intermediate row counts, as these were very expensive. Retained only essential logging to monitor spark job progress and avoid redundant computations. This change significantly improved the runtime of both jobs.
Below are the results of the model-ready data and the split train, validation, and test sets:
model_ready_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/full_model_ready_data.parquet")
model_ready_df.count()
211926
model_ready_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]Sidney_Per...| 1|
|en[SEP]V._Lakshmi...| 0|
|en[SEP]Joma_Music...| 1|
|en[SEP]Robert_Ell...| 1|
|en[SEP]New_York_D...| 0|
|en[SEP]Zeppelin_U...| 1|
|en[SEP]Melonie_Gi...| 1|
|en[SEP]RSM_US[SEP...| 1|
|en[SEP]UK_Trade_&...| 1|
|en[SEP]Arizona_He...| 1|
+--------------------+-----+
only showing top 10 rows
model_ready_df.groupBy("label").count().show()
+-----+------+
|label| count|
+-----+------+
| 1|105965|
| 0|105961|
+-----+------+
train_data_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/train_data.parquet")
train_data_df.count()
169705
train_data_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]Hervé_L._L...| 0|
|en[SEP]Elevate_Ca...| 0|
|en[SEP]Viking_(cr...| 0|
|en[SEP]Namukku_Pa...| 1|
|en[SEP]Sajid_Nadi...| 0|
|en[SEP]Endress+Ha...| 0|
|en[SEP]Gu_Su[SEP]...| 0|
|en[SEP]Anita_Dong...| 1|
|en[SEP]Glorystar[...| 0|
|en[SEP]Motorola_Q...| 1|
+--------------------+-----+
only showing top 10 rows
train_data_df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
| 1|84976|
| 0|84729|
+-----+-----+
validation_data_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/validation_data.parquet")
validation_data_df.count()
21114
validation_data_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]Oracle_App...| 0|
|en[SEP]Elie_Tahar...| 1|
|en[SEP]Jay_Richar...| 0|
|en[SEP]Friends_fr...| 0|
|en[SEP]Floridians...| 1|
|en[SEP]Yogoda_Sat...| 0|
|en[SEP]Huma_(comp...| 1|
|en[SEP]Steinway_&...| 0|
|en[SEP]Litesa_Wal...| 0|
|en[SEP]Tauț[SEP]T...| 1|
+--------------------+-----+
only showing top 10 rows
validation_data_df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
| 1|10438|
| 0|10676|
+-----+-----+
test_data_df = spark.read.parquet("/user/kevinbazira/tone_check_data_generation_pipeline/20250922T151655/test_data.parquet")
test_data_df.count()
21023
test_data_df.show(10)
+--------------------+-----+
| input|label|
+--------------------+-----+
|en[SEP]MongoDB[SE...| 0|
|en[SEP]Globant[SE...| 0|
|en[SEP]Nell_(film...| 1|
|en[SEP]Chameleon_...| 1|
|en[SEP]Signet_Arm...| 1|
|en[SEP]Samsung_El...| 0|
|en[SEP]Our_Politi...| 0|
|en[SEP]The_Myth_o...| 1|
|en[SEP]Aso_ni_San...| 1|
|en[SEP]California...| 1|
+--------------------+-----+
only showing top 10 rows
test_data_df.groupBy("label").count().show()
+-----+-----+
|label|count|
+-----+-----+
| 1|10442|
| 0|10581|
+-----+-----+Comment Actions
kevinbazira opened https://gitlab.wikimedia.org/repos/machine-learning/ml-pipelines/-/merge_requests/52
tone-check: add data generation pipeline job logic
Comment Actions
kevinbazira opened https://gitlab.wikimedia.org/repos/data-engineering/airflow-dags/-/merge_requests/1699
ml: add tone-check data generation DAG
Comment Actions
kevinbazira merged https://gitlab.wikimedia.org/repos/machine-learning/ml-pipelines/-/merge_requests/52
tone-check: add data generation pipeline job logic
Comment Actions
kevinbazira merged https://gitlab.wikimedia.org/repos/data-engineering/airflow-dags/-/merge_requests/1699
ml: add tone-check data generation DAG