|
|
|
|
|
|
|
|
|
|
|
|
|
|
from datasets import load_dataset |
|
|
from transformers import AutoTokenizer, T5ForConditionalGeneration |
|
|
from trl import SFTTrainer, SFTConfig |
|
|
|
|
|
dataset = load_dataset("mindchain/container-status-de", split="train") |
|
|
split = dataset.train_test_split(test_size=0.15, seed=42) |
|
|
|
|
|
def fmt(ex): |
|
|
return {"text": f"Status: {ex['text']}", "label": ex["label"]} |
|
|
|
|
|
train_ds = split["train"].map(fmt, remove_columns=split["train"].column_names) |
|
|
eval_ds = split["test"].map(fmt, remove_columns=split["test"].column_names) |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("google/t5gemma-2-270m") |
|
|
model = T5ForConditionalGeneration.from_pretrained("google/t5gemma-2-270m") |
|
|
|
|
|
config = SFTConfig( |
|
|
output_dir="out", |
|
|
push_to_hub=True, |
|
|
hub_model_id="mindchain/t5gemma-270m-container-status", |
|
|
num_train_epochs=5, |
|
|
per_device_train_batch_size=2, |
|
|
gradient_accumulation_steps=4, |
|
|
learning_rate=3e-4, |
|
|
logging_steps=5, |
|
|
max_length=256, |
|
|
report_to="trackio", |
|
|
) |
|
|
|
|
|
trainer = SFTTrainer(model=model, tokenizer=tokenizer, train_dataset=train_ds, eval_dataset=eval_ds, args=config) |
|
|
trainer.train() |
|
|
trainer.push_to_hub() |
|
|
print('DONE') |
|
|
|
