# 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')