Implementation:Datajuicer Data juicer Quality Classifier Utils
| Knowledge Sources | |
|---|---|
| Domains | Data_Quality, Tooling |
| Last Updated | 2026-02-14 16:00 GMT |
Overview
Concrete tool for PySpark-based quality classifier operations including Spark initialization, model management, dataset loading, tokenization, training, evaluation, and prediction provided by Data-Juicer.
Description
qc_utils is the core utility library for the quality classifier toolkit. It provides functions for every stage of the ML pipeline: init_spark configures and creates a SparkSession with tunable memory and partition settings; prepare_model loads a PipelineModel from local cache or downloads it from a remote OSS repository; load_dataset and load_datasets read JSON, JSONL, or Parquet files into PySpark DataFrames with optional labeling; shuffle randomizes DataFrame row order; export_result writes DataFrames to disk; tokenize_dataset applies sentencepiece tokenizers; get_keep_method_udf returns user-defined functions for Pareto-based or threshold-based keep decisions; train builds and fits a HashingTF plus LogisticRegression pipeline; eval computes precision, recall, and F1 on a labeled test set; and predict scores documents and adds doc_score and should_keep columns.
Usage
Use when building, training, evaluating, or running quality classifier models, as all quality classifier tools (predict.py, train.py) depend on these shared utility functions.
Code Reference
Source Location
- Repository: Datajuicer_Data_juicer
- File: data_juicer/tools/quality_classifier/qc_utils.py
Signature
def init_spark(spark_executor_memory=None, spark_driver_memory=None,
spark_executor_memoryOverhead=None):
"""Initialize a PySpark session with configurable memory settings."""
def prepare_model(model_name, model_path=DATA_JUICER_MODELS_CACHE):
"""Load a quality classifier PipelineModel from cache or download from remote."""
def load_dataset(spark, ds_path, text_key="text", only_text=False):
"""Load a single dataset (JSON, JSONL, or Parquet) into a PySpark DataFrame."""
def load_datasets(spark, ds_paths, text_key="text", label=None, only_text=True):
"""Load and union multiple datasets with optional label assignment."""
def shuffle(df):
"""Shuffle a PySpark DataFrame using a fixed random seed."""
def export_result(ds, res_path):
"""Export a DataFrame to JSON, JSONL, or Parquet format."""
def get_keep_method_udf(keep_method):
"""Return a PySpark UDF for the specified keep method ('gpt3' or 'label')."""
def tokenize_dataset(ds, tokenizer):
"""Tokenize texts using a sentencepiece tokenizer, adding a 'words' column."""
def train(output_model_path, ds, tokenizer=None):
"""Train a HashingTF + LogisticRegression pipeline and save the model."""
def eval(model_path, ds, tokenizer=None):
"""Evaluate a quality classifier model computing precision, recall, and F1."""
def predict(model, ds, tokenizer=None, keep_method="label"):
"""Score documents with doc_score and should_keep columns."""
Import
from data_juicer.tools.quality_classifier.qc_utils import (
init_spark,
prepare_model,
load_dataset,
load_datasets,
shuffle,
export_result,
get_keep_method_udf,
tokenize_dataset,
train,
eval,
predict,
)
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| spark_executor_memory | str or None | No | Spark executor memory (default: "64g") |
| spark_driver_memory | str or None | No | Spark driver memory (default: "64g") |
| model_name | str | Yes (prepare_model) | Model name ("gpt3", "chinese", "code") or path to custom model |
| ds_path / ds_paths | str or list[str] | Yes (load) | Path(s) to dataset files (JSON, JSONL, Parquet) |
| text_key | str | No | Column name holding text content (default: "text") |
| tokenizer | str or None | No | Sentencepiece model path; None uses PySpark standard Tokenizer |
| keep_method | str | No | Keep method name: "gpt3" or "label" |
| output_model_path | str | Yes (train) | Path to save trained model |
Outputs
| Name | Type | Description |
|---|---|---|
| SparkSession | pyspark.sql.SparkSession | Configured Spark session (init_spark) |
| PipelineModel | pyspark.ml.PipelineModel | Loaded or trained classifier model (prepare_model, train) |
| DataFrame | pyspark.sql.DataFrame | Loaded, tokenized, or predicted dataset |
| Files on disk | JSON/JSONL/Parquet | Exported result dataset or saved model directory |
Usage Examples
Initialize Spark and Load a Model
from data_juicer.tools.quality_classifier.qc_utils import init_spark, prepare_model
spark = init_spark(spark_executor_memory="32g", spark_driver_memory="32g")
model = prepare_model("gpt3")
Load Dataset and Predict
from data_juicer.tools.quality_classifier.qc_utils import (
load_dataset, predict, export_result
)
ds = load_dataset(spark, "./data/input.jsonl", text_key="text")
scored = predict(model, ds, tokenizer=None, keep_method="gpt3")
export_result(scored, "./data/scored_output.jsonl")
Train a Custom Model
from data_juicer.tools.quality_classifier.qc_utils import (
load_datasets, shuffle, train
)
pos = load_datasets(spark, ["pos.parquet"], label=1, only_text=True)
neg = load_datasets(spark, ["neg.parquet"], label=0, only_text=True)
ds = shuffle(pos.unionAll(neg))
train("./models/my_quality_model", ds, tokenizer=None)