def create_pseudo_labeled_data(args, infer_input, infer_output, eval_result, id2label, next_data_dir):
    """Create pseudeo labeled data for the next self-training iteration."""

    dataset = datasets.concatenate_datasets([infer_input, infer_output], axis=1)

    if args.do_filter_by_confidence:
        dataset = dataset.filter(lambda example: example["probability"] > args.confidence_threshold)

    if args.do_filter_by_val_performance:
        assert eval_result >= 0.0 and eval_result <= 1.0
        num_selected_rows = int(eval_result * len(dataset))
        print(num_selected_rows)
        dataset = dataset.sort("probability", reverse=True)
        dataset = dataset.select(range(num_selected_rows))

    dataset = dataset.remove_columns(["label", "probability"])
    dataset = dataset.rename_column("prediction", "label")
    dataset = dataset.map(lambda example: {"label": id2label[example["label"]]})
    dataset = dataset.shuffle(seed=args.seed)

    pseudo_labeled_data_file = os.path.join(next_data_dir, f"train_pseudo.{args.data_file_extension}")
    if args.data_file_extension == "csv":
        dataset.to_csv(pseudo_labeled_data_file, index=False)
    else:
        dataset.to_json(pseudo_labeled_data_file)
Create pseudeo labeled data for the next self-training iteration.
dpkd/transformers/examples/research_projects/self-training-text-classification/selftraining.py