tasknet is an interface between Huggingface datasets and Huggingface transformers Trainer.
Tasknet should work with all recent versions of Transformers.
pip install tasknet
Each task template has fields that should be matched with specific dataset columns. Classification has two text fields s1,s2, and a label y. Pass a dataset to a template, and fill in the mapping between the template fields and the dataset columns to instantiate a task.
import tasknet as tn; from datasets import load_dataset
rte = tn.Classification(
dataset=load_dataset("glue", "rte"),
s1="sentence1", s2="sentence2", y="label") #s2 is optional for classification, used to represent text pairs
# See AutoTask for shorter code
class hparams:
model_name = 'tasksource/ModernBERT-base-nli' # better performance for most tasks
learning_rate = 3e-5 # see hf.co/docs/transformers/en/main_classes/trainer#transformers.TrainingArguments
model, trainer = tn.Model_Trainer(tasks=[rte], args=hparams)
trainer.train(), trainer.evaluate()
p = trainer.pipeline()
p([{'text':'premise here','text_pair': 'hypothesis here'}]) # HuggingFace pipeline for inferenceTasknet is multitask by design. model.task_models_list contains one model per task, with a shared encoder.
class hparams:
model_name = 'jhu-clsp/ettin-encoder-150m'
pack_encoder = True # padding-free fused kernels from packed-encoders
bf16 = Truepack_encoder installs packed-encoders on the shared encoder (pip install 'packed-encoders[cuda]', torch 2.11). It applies to ModernBERT-family encoders (ModernBERT, Ettin, mmBERT) and replaces the task embedding (add_cls), which it disables.
For UltraBERT, padding_free = True runs the encoder on the real tokens only (FlashAttention-2 varlen, pip install flash-attn): 1.4x faster training steps on GLUE pairs, same outputs as the padded forward up to bf16 noise.
Weights load in the checkpoint's stored dtype (torch_dtype overrides). fp16 weights, such as deberta-v3's, are converted to bf16: AdamW on fp16 weights is NaN from the first update. bf16 weights train with tasknet.optim.AdamWSR, AdamW with stochastic rounding, unless you set another optimizer: plain AdamW on bf16 weights rounds small updates away (deberta-v3-base, MNLI, 1000 steps: 0.47-0.55 accuracy plain, 0.863-0.867 with stochastic rounding, 0.854-0.861 with fp32 weights and bf16 autocast, at 35% lower peak memory and a slower optimizer step). For fp32 weights pass torch_dtype = torch.float32. Why: docs/16bit_training.md. Training stops with a FloatingPointError as soon as it diverges: non-finite gradients or loss, or a loss above 10x the first step's loss for 3 optimizer steps in a row (divergence_guard_kwargs = dict(factor=10, patience=3); divergence_guard = False disables it).
tasknet relies on task templates to avoid boilerplate codes. The task templates correspond to Transformers AutoClasses:
Classification(s1, s2, y)TokenClassification(tokens, labels)(tokens and labels are lists of words and assigned labels)MultipleChoice(s1, choices)(s1 is a prompt/question, choices is a list of texts, y is the index of the correct choice)QuestionAnswering(context, question, answers)(extractive QA in the SQuAD format; empty answers mean unanswerable)SpanTagging(question, context, spans)(query-conditioned extraction: character spans of the context answering the question, any number, none for no answer)Seq2SeqLM(experimental support)
The task templates follow the same interface. They implement preprocess_function, a data collator and compute_metrics.
Look at tasks.py and use existing templates as a starting point to implement a custom task template.
You can also leverage tasksource with tn.AutoTask and have one-line access to 600+ datasets, see implemented tasks.
rte = tn.AutoTask("glue/rte", max_rows=5000)AutoTask guesses a template based on the dataset structure. It also accepts a dataset as input, if it fits the template (e.g. after tasksource custom preprocessing).
tn.Classification(dataset, max_rows=5000, max_rows_eval=500, oversampling=2)You can balance multiple datasets with max_rows and oversampling. max_rows is the maximal number of examples. If a dataset has less than max_rows, it will be oversampled at most oversampling times.
Compare target-only training with source+target training at the same optimization budget:
result = tn.transfer(
sources=["glue/mrpc", "glue/qqp"],
targets=["glue/rte", "glue/cola"],
model="answerdotai/ModernBERT-base",
nrows=1000, steps=200, seeds=[0, 1],
output_dir="results/study",
)
print(result.matrix()) # mean treatment minus baseline, by source and targetThe experiment API uses nrows; existing task constructors use max_rows.
Start with classification-to-classification or MCQ-to-MCQ comparisons, which
reuse their existing templates and heads. Mixed task types are exploratory:
the source fraction controls batches rather than examples or tokens.
See the transfer protocol and examples for controls, multi-source
studies, UltraBERT classification and MCQ, speed controls, CLI configuration, and JSON/CSV results.
Transfer reuses tokenization and initial target weights, skips unused source evaluation
processing, and honors a separate evaluation batch size.
tasknet transfer --config study.yaml
tasknet report results/study
tasknet list-tasksMinimal-ish example:
https://colab.research.google.com/drive/15Xf4Bgs3itUmok7XlAK6EEquNbvjD9BD?usp=sharing
More complex example, where tasknet was scaled to 600 tasks:
https://colab.research.google.com/drive/1iB4Oxl9_B5W3ZDzXoWJN-olUbqLBxgQS?usp=sharing
This code uses some part of the examples of the transformers library and some code from multitask-learning-transformers.
You can request features on github or reach me at damien.sileo@inria.fr
@misc{sileod22-tasknet,
author = {Sileo, Damien},
doi = {10.5281/zenodo.561225781},
month = {11},
title = {{tasknet, multitask interface between Trainer and datasets}},
url = {https://github.com/sileod/tasknet},
version = {1.5.0},
year = {2022}}