Views
No views yet
🎯 TL;DR: State-of-the-art paired encoder and decoder models (17M-1B params) trained identically for fair comparison with open data. Encoders beat ModernBERT. Decoders beat Llama 3.2/SmolLM2.
1pip install torch>=1.9.0
2# until the new pip release, install from main to use decoders (transformers>=4.54.X will contain it)
3# encoders work with transformers>=4.48.0
4pip install git+https://github.com/huggingface/transformers.git1from transformers import AutoTokenizer, AutoModel
2
3tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-150m")
4model = AutoModel.from_pretrained("jhu-clsp/ettin-encoder-150m")1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-150m")
4model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-150m")| Size | Model | Parameters | Best For | Download |
|---|---|---|---|---|
| XXS | ettin-encoder-17m | 17M | Mobile/Edge devices | |
| XS | ettin-encoder-32m | 32M | Fast inference | |
| Small | ettin-encoder-68m | 68M | Balanced performance | |
| Base | ettin-encoder-150m | 150M | Standard use cases | |
| Large | ettin-encoder-400m | 400M | High accuracy needs | |
| XL | ettin-encoder-1b | 1B | Best performance |
| Size | Model | Parameters | Best For | Download |
|---|---|---|---|---|
| XXS | ettin-decoder-17m | 17M | Lightweight generation | |
| XS | ettin-decoder-32m | 32M | Quick prototyping | |
| Small | ettin-decoder-68m | 68M | Efficient generation | |
| Base | ettin-decoder-150m | 150M | Standard generation | |
| Large | ettin-decoder-400m | 400M | Quality generation | |
| XL | ettin-decoder-1b | 1B | Best generation |
AutoModel or AutoModelForMaskedLM:| Size | Model | Parameters | Description | Download |
|---|---|---|---|---|
| XXS | ettin-encoder-from-decoder-17m | 17M | Decoder → MLM continued training | |
| XS | ettin-encoder-from-decoder-32m | 32M | Decoder → MLM continued training | |
| Small | ettin-encoder-from-decoder-68m | 68M | Decoder → MLM continued training | |
| Base | ettin-encoder-from-decoder-150m | 150M | Decoder → MLM continued training | |
| Large | ettin-encoder-from-decoder-400m | 400M | Decoder → MLM continued training | |
| XL | ettin-encoder-from-decoder-1b | 1B | Decoder → MLM continued training |
AutoModelForCausalLM:| Size | Model | Parameters | Description | Download |
|---|---|---|---|---|
| XXS | ettin-decoder-from-encoder-17m | 17M | Encoder → CLM continued training | |
| XS | ettin-decoder-from-encoder-32m | 32M | Encoder → CLM continued training | |
| Small | ettin-decoder-from-encoder-68m | 68M | Encoder → CLM continued training | |
| Base | ettin-decoder-from-encoder-150m | 150M | Encoder → CLM continued training | |
| Large | ettin-decoder-from-encoder-400m | 400M | Encoder → CLM continued training | |
| XL | ettin-decoder-from-encoder-1b | 1B | Encoder → CLM continued training |
1# Encoder-from-decoder: Load as encoder
2from transformers import AutoTokenizer, AutoModel
3tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-from-decoder-150m")
4model = AutoModel.from_pretrained("jhu-clsp/ettin-encoder-from-decoder-150m")
5
6# Decoder-from-encoder: Load as decoder
7from transformers import AutoTokenizer, AutoModelForCausalLM
8tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-from-encoder-150m")
9model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-from-encoder-150m")step{number} - Pretraining phase checkpoints (e.g., step599525, step596528)ext{number} - Extension/mid-training phase checkpoints (e.g., ext1000, ext2000)decay{number} - Decay phase checkpoints (e.g., decay100, decay500)1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3# Load a specific pretraining checkpoint
4model = AutoModelForCausalLM.from_pretrained(
5 "jhu-clsp/ettin-decoder-400m",
6 revision="step590532" # Specific checkpoint tag
7)
8
9# Load an extension phase checkpoint
10model = AutoModelForCausalLM.from_pretrained(
11 "jhu-clsp/ettin-decoder-400m",
12 revision="ext1000"
13)
14
15# Load a decay phase checkpoint
16model = AutoModelForCausalLM.from_pretrained(
17 "jhu-clsp/ettin-decoder-400m",
18 revision="decay100"
19)| Parameter | 17M | 32M | 68M | 150M | 400M | 1B |
|---|---|---|---|---|---|---|
| Layers | 7 | 10 | 19 | 22 | 28 | 28 |
| Hidden Size | 256 | 384 | 512 | 768 | 1024 | 1792 |
| Intermediate Size | 384 | 576 | 768 | 1152 | 2624 | 3840 |
| Attention Heads | 4 | 6 | 8 | 12 | 16 | 28 |
1from transformers import AutoTokenizer, AutoModelForMaskedLM
2import torch
3
4# Load MLM model
5tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-150m")
6model = AutoModelForMaskedLM.from_pretrained("jhu-clsp/ettin-encoder-150m")
7
8def predict_masked_token(text):
9 inputs = tokenizer(text, return_tensors="pt")
10 with torch.no_grad():
11 outputs = model(**inputs)
12
13 # Get predictions for [MASK] tokens
14 mask_indices = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)
15 predictions = outputs.logits[mask_indices]
16
17 # Get top 5 predictions
18 top_tokens = torch.topk(predictions, 5, dim=-1)
19 return [tokenizer.decode(token) for token in top_tokens.indices[0]]
20
21# Example
22masked_text = "The capital of France is [MASK]."
23predictions = predict_masked_token(masked_text)
24print(f"Predictions: {predictions}")1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4# Load model and tokenizer
5tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-150m")
6model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-150m")
7
8# Set pad token if needed
9if tokenizer.pad_token is None:
10 tokenizer.pad_token = tokenizer.eos_token
11
12def generate_text(prompt, max_length=100, temperature=0.7):
13 inputs = tokenizer(prompt, return_tensors="pt")
14
15 with torch.no_grad():
16 outputs = model.generate(
17 inputs.input_ids,
18 max_length=max_length,
19 temperature=temperature,
20 do_sample=True,
21 pad_token_id=tokenizer.eos_token_id,
22 num_return_sequences=1
23 )
24
25 return tokenizer.decode(outputs[0], skip_special_tokens=True)
26
27# Example usage
28prompt = "The future of artificial intelligence is"
29generated = generate_text(prompt)
30print(generated)1import argparse
2
3from datasets import load_dataset
4from sentence_transformers import (
5 SentenceTransformer,
6 SentenceTransformerTrainer,
7 SentenceTransformerTrainingArguments,
8)
9from sentence_transformers.evaluation import TripletEvaluator
10from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
11from sentence_transformers.training_args import BatchSamplers
12
13def main():
14 # parse the lr & model name
15 parser = argparse.ArgumentParser()
16 parser.add_argument("--lr", type=float, default=8e-5)
17 parser.add_argument("--model_name", type=str, default="jhu-clsp/ettin-encoder-150m")
18 args = parser.parse_args()
19 lr = args.lr
20 model_name = args.model_name
21 model_shortname = model_name.split("/")[-1]
22
23 # 1. Load a model to finetune
24 model = SentenceTransformer(model_name)
25
26 # 2. Load a dataset to finetune on
27 dataset = load_dataset(
28 "sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1",
29 "triplet-hard",
30 split="train",
31 )
32 dataset_dict = dataset.train_test_split(test_size=1_000, seed=12)
33 train_dataset = dataset_dict["train"].select(range(1_250_000))
34 eval_dataset = dataset_dict["test"]
35
36 # 3. Define a loss function
37 loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16) # Increase mini_batch_size if you have enough VRAM
38
39 run_name = f"{model_shortname}-DPR-{lr}"
40 # 4. (Optional) Specify training arguments
41 args = SentenceTransformerTrainingArguments(
42 # Required parameter:
43 output_dir=f"output/{model_shortname}/{run_name}",
44 # Optional training parameters:
45 num_train_epochs=1,
46 per_device_train_batch_size=512,
47 per_device_eval_batch_size=512,
48 warmup_ratio=0.05,
49 fp16=False, # Set to False if GPU can't handle FP16
50 bf16=True, # Set to True if GPU supports BF16
51 batch_sampler=BatchSamplers.NO_DUPLICATES, # (Cached)MultipleNegativesRankingLoss benefits from no duplicates
52 learning_rate=lr,
53 # Optional tracking/debugging parameters:
54 save_strategy="steps",
55 save_steps=500,
56 save_total_limit=2,
57 logging_steps=500,
58 run_name=run_name, # Used in `wandb`, `tensorboard`, `neptune`, etc. if installed
59 )
60
61 # 5. (Optional) Create an evaluator & evaluate the base model
62 dev_evaluator = TripletEvaluator(
63 anchors=eval_dataset["query"],
64 positives=eval_dataset["positive"],
65 negatives=eval_dataset["negative"],
66 name="msmarco-co-condenser-dev",
67 )
68 dev_evaluator(model)
69
70 # 6. Create a trainer & train
71 trainer = SentenceTransformerTrainer(
72 model=model,
73 args=args,
74 train_dataset=train_dataset,
75 eval_dataset=eval_dataset,
76 loss=loss,
77 evaluator=dev_evaluator,
78 )
79 trainer.train()
80
81 # 7. (Optional) Evaluate the trained model on the evaluator after training
82 dev_evaluator(model)
83
84 # 8. Save the model
85 model.save_pretrained(f"output/{model_shortname}/{run_name}/final")
86
87 # 9. (Optional) Push it to the Hugging Face Hub
88 model.push_to_hub(run_name, private=False)
89
90if __name__ == "__main__":
91 main()1from datasets import load_dataset
2from pylate import losses, models, utils
3from sentence_transformers import (
4 SentenceTransformerTrainer,
5 SentenceTransformerTrainingArguments,
6)
7
8def main():
9 # Load the datasets required for knowledge distillation (train, queries, documents)
10 train = load_dataset(
11 path="lightonai/ms-marco-en-bge",
12 name="train",
13 )
14
15 queries = load_dataset(
16 path="lightonai/ms-marco-en-bge",
17 name="queries",
18 )
19
20 documents = load_dataset(
21 path="lightonai/ms-marco-en-bge",
22 name="documents",
23 )
24
25 # Set the transformation to load the documents/queries texts using the corresponding ids on the fly
26 train.set_transform(
27 utils.KDProcessing(queries=queries, documents=documents).transform,
28 )
29
30 # Define the base model, training parameters, and output directory
31 num_train_epochs = 1
32 lr = 8e-5
33 batch_size = 16
34 accum_steps = 1
35 model_name = "jhu-clsp/ettin-encoder-150m"
36 model_shortname = model_name.split("/")[-1]
37
38 # Set the run name for logging and output directory
39 run_name = f"{model_shortname}-colbert-KD-{lr}"
40 output_dir = f"output/{model_shortname}/{run_name}"
41
42 # Initialize the ColBERT model from the base model
43 model = models.ColBERT(model_name_or_path=model_name)
44
45 # Configure the training arguments (e.g., epochs, batch size, learning rate)
46 args = SentenceTransformerTrainingArguments(
47 output_dir=output_dir,
48 num_train_epochs=num_train_epochs,
49 per_device_train_batch_size=batch_size,
50 fp16=False, # Set to False if you get an error that your GPU can't run on FP16
51 bf16=True, # Set to True if you have a GPU that supports BF16
52 run_name=run_name,
53 logging_steps=10,
54 learning_rate=lr,
55 gradient_accumulation_steps=accum_steps,
56 warmup_ratio=0.05,
57 )
58
59 # Use the Distillation loss function for training
60 train_loss = losses.Distillation(model=model)
61
62 # Initialize the trainer
63 trainer = SentenceTransformerTrainer(
64 model=model,
65 args=args,
66 train_dataset=train,
67 loss=train_loss,
68 data_collator=utils.ColBERTCollator(tokenize_fn=model.tokenize),
69 )
70
71 # Start the training process
72 trainer.train()
73
74 model.save_pretrained(f"{output_dir}/final")
75
76if __name__ == "__main__":
77 main()
781import logging
2
3from datasets import load_dataset
4
5from sentence_transformers import (
6 SparseEncoder,
7 SparseEncoderModelCardData,
8 SparseEncoderTrainer,
9 SparseEncoderTrainingArguments,
10)
11from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
12from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
13from sentence_transformers.training_args import BatchSamplers
14
15logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
16
17# 1. Load a model to finetune with 2. (Optional) model card data
18model = SparseEncoder(
19 "jhu-clsp/ettin-encoder-150m",
20 model_card_data=SparseEncoderModelCardData(
21 language="en",
22 license="apache-2.0",
23 )
24)
25
26# 3. Load a dataset to finetune on
27full_dataset = load_dataset("sentence-transformers/natural-questions", split="train").select(range(100_000))
28dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
29train_dataset = dataset_dict["train"]
30eval_dataset = dataset_dict["test"]
31
32# 4. Define a loss function
33loss = SpladeLoss(
34 model=model,
35 loss=SparseMultipleNegativesRankingLoss(model=model),
36 query_regularizer_weight=5e-5,
37 document_regularizer_weight=3e-5,
38)
39
40# 5. (Optional) Specify training arguments
41run_name = "splade-distilbert-base-uncased-nq"
42args = SparseEncoderTrainingArguments(
43 # Required parameter:
44 output_dir=f"models/{run_name}",
45 # Optional training parameters:
46 num_train_epochs=1,
47 per_device_train_batch_size=16,
48 per_device_eval_batch_size=16,
49 learning_rate=2e-5,
50 warmup_ratio=0.1,
51 fp16=True, # Set to False if you get an error that your GPU can't run on FP16
52 bf16=False, # Set to True if you have a GPU that supports BF16
53 batch_sampler=BatchSamplers.NO_DUPLICATES, # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
54 # Optional tracking/debugging parameters:
55 eval_strategy="steps",
56 eval_steps=1000,
57 save_strategy="steps",
58 save_steps=1000,
59 save_total_limit=2,
60 logging_steps=200,
61 run_name=run_name, # Will be used in W&B if `wandb` is installed
62)
63
64# 6. (Optional) Create an evaluator & evaluate the base model
65dev_evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"], batch_size=16)
66
67# 7. Create a trainer & train
68trainer = SparseEncoderTrainer(
69 model=model,
70 args=args,
71 train_dataset=train_dataset,
72 eval_dataset=eval_dataset,
73 loss=loss,
74 evaluator=dev_evaluator,
75)
76trainer.train()
77
78# 8. Evaluate the model performance again after training
79dev_evaluator(model)
80
81# 9. Save the trained model
82model.save_pretrained(f"models/{run_name}/final")
83
84# 10. (Optional) Push it to the Hugging Face Hub
85model.push_to_hub(run_name)
861import logging
2import traceback
3
4import torch
5from datasets import load_dataset
6
7from sentence_transformers import SentenceTransformer
8from sentence_transformers.cross_encoder import (
9 CrossEncoder,
10 CrossEncoderModelCardData,
11 CrossEncoderTrainer,
12 CrossEncoderTrainingArguments,
13)
14from sentence_transformers.cross_encoder.evaluation import (
15 CrossEncoderNanoBEIREvaluator,
16 CrossEncoderRerankingEvaluator,
17)
18from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
19from sentence_transformers.evaluation import SequentialEvaluator
20from sentence_transformers.util import mine_hard_negatives
21
22# Set the log level to INFO to get more information
23logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
24
25
26def main():
27 model_name = "jhu-clsp/ettin-encoder-150m"
28
29 train_batch_size = 64
30 num_epochs = 1
31 num_hard_negatives = 5 # How many hard negatives should be mined for each question-answer pair
32
33 # 1a. Load a model to finetune with 1b. (Optional) model card data
34 model = CrossEncoder(
35 model_name,
36 model_card_data=CrossEncoderModelCardData(
37 language="en",
38 license="apache-2.0",
39 ),
40 )
41 print("Model max length:", model.max_length)
42 print("Model num labels:", model.num_labels)
43
44 # 2a. Load the GooAQ dataset: https://huggingface.co/datasets/sentence-transformers/gooaq
45 logging.info("Read the gooaq training dataset")
46 full_dataset = load_dataset("sentence-transformers/gooaq", split="train").select(range(100_000))
47 dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
48 train_dataset = dataset_dict["train"]
49 eval_dataset = dataset_dict["test"]
50 logging.info(train_dataset)
51 logging.info(eval_dataset)
52
53 # 2b. Modify our training dataset to include hard negatives using a very efficient embedding model
54 embedding_model = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1", device="cpu")
55 hard_train_dataset = mine_hard_negatives(
56 train_dataset,
57 embedding_model,
58 num_negatives=num_hard_negatives, # How many negatives per question-answer pair
59 margin=0, # Similarity between query and negative samples should be x lower than query-positive similarity
60 range_min=0, # Skip the x most similar samples
61 range_max=100, # Consider only the x most similar samples
62 sampling_strategy="top", # Sample the top negatives from the range
63 batch_size=4096, # Use a batch size of 4096 for the embedding model
64 output_format="labeled-pair", # The output format is (query, passage, label), as required by BinaryCrossEntropyLoss
65 use_faiss=True,
66 )
67 logging.info(hard_train_dataset)
68
69 # 2c. (Optionally) Save the hard training dataset to disk
70 # hard_train_dataset.save_to_disk("gooaq-hard-train")
71 # Load again with:
72 # hard_train_dataset = load_from_disk("gooaq-hard-train")
73
74 # 3. Define our training loss.
75 # pos_weight is recommended to be set as the ratio between positives to negatives, a.k.a. `num_hard_negatives`
76 loss = BinaryCrossEntropyLoss(model=model, pos_weight=torch.tensor(num_hard_negatives))
77
78 # 4a. Define evaluators. We use the CrossEncoderNanoBEIREvaluator, which is a light-weight evaluator for English reranking
79 nano_beir_evaluator = CrossEncoderNanoBEIREvaluator(
80 dataset_names=["msmarco", "nfcorpus", "nq"],
81 batch_size=train_batch_size,
82 )
83
84 # 4b. Define a reranking evaluator by mining hard negatives given query-answer pairs
85 # We include the positive answer in the list of negatives, so the evaluator can use the performance of the
86 # embedding model as a baseline.
87 hard_eval_dataset = mine_hard_negatives(
88 eval_dataset,
89 embedding_model,
90 corpus=full_dataset["answer"], # Use the full dataset as the corpus
91 num_negatives=30, # How many documents to rerank
92 batch_size=4096,
93 include_positives=True,
94 output_format="n-tuple",
95 use_faiss=True,
96 )
97 logging.info(hard_eval_dataset)
98 reranking_evaluator = CrossEncoderRerankingEvaluator(
99 samples=[
100 {
101 "query": sample["question"],
102 "positive": [sample["answer"]],
103 "documents": [sample[column_name] for column_name in hard_eval_dataset.column_names[2:]],
104 }
105 for sample in hard_eval_dataset
106 ],
107 batch_size=train_batch_size,
108 name="gooaq-dev",
109 # Realistic setting: only rerank the positives that the retriever found
110 # Set to True to rerank *all* positives
111 always_rerank_positives=False,
112 )
113
114 # 4c. Combine the evaluators & run the base model on them
115 evaluator = SequentialEvaluator([reranking_evaluator, nano_beir_evaluator])
116 evaluator(model)
117
118 # 5. Define the training arguments
119 short_model_name = model_name if "/" not in model_name else model_name.split("/")[-1]
120 run_name = f"reranker-{short_model_name}-gooaq-bce"
121 args = CrossEncoderTrainingArguments(
122 # Required parameter:
123 output_dir=f"models/{run_name}",
124 # Optional training parameters:
125 num_train_epochs=num_epochs,
126 per_device_train_batch_size=train_batch_size,
127 per_device_eval_batch_size=train_batch_size,
128 learning_rate=2e-5,
129 warmup_ratio=0.1,
130 fp16=False, # Set to False if you get an error that your GPU can't run on FP16
131 bf16=True, # Set to True if you have a GPU that supports BF16
132 dataloader_num_workers=4,
133 load_best_model_at_end=True,
134 metric_for_best_model="eval_gooaq-dev_ndcg@10",
135 # Optional tracking/debugging parameters:
136 eval_strategy="steps",
137 eval_steps=1000,
138 save_strategy="steps",
139 save_steps=1000,
140 save_total_limit=2,
141 logging_steps=200,
142 logging_first_step=True,
143 run_name=run_name, # Will be used in W&B if `wandb` is installed
144 seed=12,
145 )
146
147 # 6. Create the trainer & start training
148 trainer = CrossEncoderTrainer(
149 model=model,
150 args=args,
151 train_dataset=hard_train_dataset,
152 loss=loss,
153 evaluator=evaluator,
154 )
155 trainer.train()
156
157 # 7. Evaluate the final model, useful to include these in the model card
158 evaluator(model)
159
160 # 8. Save the final model
161 final_output_dir = f"models/{run_name}/final"
162 model.save_pretrained(final_output_dir)
163
164 # 9. (Optional) save the model to the Hugging Face Hub!
165 # It is recommended to run `huggingface-cli login` to log into your Hugging Face account first
166 try:
167 model.push_to_hub(run_name)
168 except Exception:
169 logging.error(
170 f"Error uploading model to the Hugging Face Hub:\n{traceback.format_exc()}To upload it manually, you can run "
171 f"`huggingface-cli login`, followed by loading the model using `model = CrossEncoder({final_output_dir!r})` "
172 f"and saving it using `model.push_to_hub('{run_name}')`."
173 )
174
175
176if __name__ == "__main__":
177 main()
1781python trl/scripts/sft.py \
2 --model_name_or_path jhu-clsp/ettin-decoder-17m \
3 --dataset_name trl-lib/Capybara \
4 --learning_rate 2.0e-5 \
5 --num_train_epochs 1 \
6 --packing \
7 --per_device_train_batch_size 2 \
8 --gradient_accumulation_steps 8 \
9 --gradient_checkpointing \
10 --eos_token '<|im_end|>' \
11 --eval_strategy steps \
12 --eval_steps 100 \
13 --output_dir ettin-decoder-17m \
14 --push_to_hub1python trl/scripts/sft.py \
2 --model_name_or_path jhu-clsp/ettin-decoder-17m \
3 --dataset_name trl-lib/Capybara \
4 --learning_rate 2.0e-4 \
5 --num_train_epochs 1 \
6 --packing \
7 --per_device_train_batch_size 2 \
8 --gradient_accumulation_steps 8 \
9 --gradient_checkpointing \
10 --eos_token '<|im_end|>' \
11 --eval_strategy steps \
12 --eval_steps 100 \
13 --use_peft \
14 --lora_r 32 \
15 --lora_alpha 16 \
16 --output_dir ettin-decoder-17m \
17 --push_to_hubsft.py:1import argparse
2
3from datasets import load_dataset
4from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
5from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
6
7from trl import (
8 ModelConfig,
9 ScriptArguments,
10 SFTConfig,
11 SFTTrainer,
12 TrlParser,
13 clone_chat_template,
14 get_kbit_device_map,
15 get_peft_config,
16 get_quantization_config,
17)
18
19
20def main(script_args, training_args, model_args):
21 ################
22 # Model init kwargs & Tokenizer
23 ################
24 quantization_config = get_quantization_config(model_args)
25 model_kwargs = dict(
26 revision=model_args.model_revision,
27 trust_remote_code=model_args.trust_remote_code,
28 attn_implementation=model_args.attn_implementation,
29 torch_dtype=model_args.torch_dtype,
30 use_cache=False if training_args.gradient_checkpointing else True,
31 device_map=get_kbit_device_map() if quantization_config is not None else None,
32 quantization_config=quantization_config,
33 )
34
35 # Create model
36 config = AutoConfig.from_pretrained(model_args.model_name_or_path)
37 valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values()
38
39 if config.architectures and any(arch in valid_image_text_architectures for arch in config.architectures):
40 from transformers import AutoModelForImageTextToText
41
42 model_kwargs.pop("use_cache", None) # Image models do not support cache
43 model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs)
44 else:
45 model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs)
46
47 # Create tokenizer
48 tokenizer = AutoTokenizer.from_pretrained(
49 model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, use_fast=True
50 )
51
52 # Set default chat template if needed
53 if tokenizer.chat_template is None:
54 # TODO: source should be passed as an argument
55 model, tokenizer = clone_chat_template(model, tokenizer, "Qwen/Qwen3-0.6B")
56
57 ################
58 # Dataset
59 ################
60 dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)
61
62 ################
63 # Training
64 ################
65 trainer = SFTTrainer(
66 model=model,
67 args=training_args,
68 train_dataset=dataset[script_args.dataset_train_split],
69 eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None,
70 processing_class=tokenizer,
71 peft_config=get_peft_config(model_args),
72 )
73
74 trainer.train()
75
76 # Save and push to hub
77 trainer.save_model(training_args.output_dir)
78 if training_args.push_to_hub:
79 trainer.push_to_hub(dataset_name=script_args.dataset_name)
80
81
82def make_parser(subparsers: argparse._SubParsersAction = None):
83 dataclass_types = (ScriptArguments, SFTConfig, ModelConfig)
84 if subparsers is not None:
85 parser = subparsers.add_parser("sft", help="Run the SFT training script", dataclass_types=dataclass_types)
86 else:
87 parser = TrlParser(dataclass_types)
88 return parser
89
90
91if __name__ == "__main__":
92 parser = make_parser()
93 # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments.
94 # To ensure that their parsing does not interfere with the script arguments, parse the arguments with
95 # `return_remaining_strings=True`, then ignore the remaining strings.
96 script_args, training_args, model_args, _ = parser.parse_args_and_config(return_remaining_strings=True)
97 main(script_args, training_args, model_args)
981@misc{weller2025seqvsseqopen,
2 title={Seq vs Seq: An Open Suite of Paired Encoders and Decoders},
3 author={Orion Weller and Kathryn Ricci and Marc Marone and Antoine Chaffin and Dawn Lawrie and Benjamin Van Durme},
4 year={2025},
5 eprint={2507.11412},
6 archivePrefix={arXiv},
7 primaryClass={cs.CL},
8 url={https://arxiv.org/abs/2507.11412},
9}