Spaces:
Runtime error
Runtime error
File size: 1,295 Bytes
df1dd29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | # train_summarizer.py
from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments
from datasets import load_dataset
# Load dataset (must have 'input_text' and 'summary' columns)
dataset = load_dataset('csv', data_files='summarization_data.csv')
# Tokenizer
tokenizer = T5Tokenizer.from_pretrained('t5-small')
# Preprocessing
def preprocess(example):
input_enc = tokenizer("summarize: " + example['input_text'], truncation=True, padding='max_length', max_length=512)
target_enc = tokenizer(example['summary'], truncation=True, padding='max_length', max_length=128)
return {
'input_ids': input_enc['input_ids'],
'attention_mask': input_enc['attention_mask'],
'labels': target_enc['input_ids']
}
tokenized = dataset.map(preprocess)
# Load model
model = T5ForConditionalGeneration.from_pretrained('t5-small')
# Training config
args = TrainingArguments(
output_dir="./t5_summarizer",
per_device_train_batch_size=4,
num_train_epochs=3,
logging_dir="./logs",
logging_steps=10,
save_strategy="no"
)
# Train
trainer = Trainer(model=model, args=args, train_dataset=tokenized['train'])
trainer.train()
# Save model
model.save_pretrained('./t5_summarizer')
tokenizer.save_pretrained('./t5_summarizer')
|