def METHOD_NAME(dataset_token_classification) -> None:
    trainer = ArgillaTrainer(name=dataset_token_classification, model=MODEL, framework=FRAMEWORK, train_size=0.8)
    trainer.update_config(max_steps=1, num_train_epochs=1)
    output_dir = f"tmp_{FRAMEWORK}_train_tokencat"
    train_with_cleanup(trainer, output_dir)
    record = trainer.predict("This is a text", as_argilla_records=True)
    assert isinstance(record, TokenClassificationRecord)
    not_record = trainer.predict("This is a text", as_argilla_records=False)
    assert not isinstance(not_record, TokenClassificationRecord)
    train_with_cleanup(trainer, output_dir, train=False)