safetensors t5 region:us
---
{}
---
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments

# 1. Tải dữ liệu
dataset = load_dataset("truongpdd/vietnews-dataset")
train_dataset = dataset["train"]
val_dataset = dataset["validation"]

# 2. Tải mô hình và tokenizer
model_name = "t5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 3. Tiền xử lý dữ liệu
def preprocess_data(examples):
    inputs = ["summarize: " + article for article in examples["article"]]
    model_inputs = tokenizer(inputs, max_length=512, truncation=True)
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(examples["summary"], max_length=150, truncation=True)
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_train = train_dataset.map(preprocess_data, batched=True)
tokenized_val = val_dataset.map(preprocess_data, batched=True)

# 4. Thiết lập tham số huấn luyện
training_args = Seq2SeqTrainingArguments(
    output_dir="./vietnews-summarization",
    evaluation_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=3,
    predict_with_generate=True,
    logging_dir="./logs",
    logging_steps=100,
    push_to_hub=False,
)

# 5. Khởi tạo và huấn luyện
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    tokenizer=tokenizer,
)
trainer.train()

# 6. Lưu mô hình
model.save_pretrained("./vietnews-summarization-model")
tokenizer.save_pretrained("./vietnews-summarization-model")

# 7. Kiểm tra mô hình với dữ liệu mẫu
article = """
VKS cáo buộc ông Nguyễn Thế Hiệp có sai phạm trong vụ cháy gần Bệnh viện Nhi trung ương khiến 2 người chết, thiệt hại 1,9 tỷ đồng song bị cáo khẳng định vô tội...
"""
inputs = tokenizer("summarize: " + article, return_tensors="pt", max_length=512, truncation=True)
summary_ids = model.generate(inputs["input_ids"], max_length=150, min_length=30, length_penalty=2.0, num_beams=4)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
print("Bản tóm tắt:", summary)