CoolFace
Apppublic

AKALYAS/RNN_Task

sourceHugging Faceupdated 1y agoView on Hugging Face
1likes
trained_model.py55 linesDownload Raw Back to root
1import pandas as pd2from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments3from datasets import Dataset4 5df = pd.read_csv("news_summary_more.csv", encoding="latin1")6df = df[["text", "headlines"]]  7df = df.rename(columns={"text": "article", "headlines": "summary"})8df = df.dropna()9 10df = df.sample(5000, random_state=42)11 12model_name = "t5-small"13tokenizer = T5Tokenizer.from_pretrained(model_name)14 15def preprocess(example):16    inputs = ["summarize: " + art for art in example["article"]]17    model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")18 19    labels = tokenizer(example["summary"], max_length=128, truncation=True, padding="max_length")20    model_inputs["labels"] = labels["input_ids"]21    return model_inputs22 23dataset = Dataset.from_pandas(df)24dataset = dataset.map(preprocess, batched=True, remove_columns=["article", "summary"])25 26train_test = dataset.train_test_split(test_size=0.1)27train_dataset = train_test["train"]28eval_dataset = train_test["test"]29 30model = T5ForConditionalGeneration.from_pretrained(model_name)31 32training_args = TrainingArguments(33    output_dir="./results",34    per_device_train_batch_size=4,35    per_device_eval_batch_size=4,36    num_train_epochs=1,   37    weight_decay=0.01,38    save_total_limit=1,39    logging_dir="./logs",40    logging_steps=5041)42 43trainer = Trainer(44    model=model,45    args=training_args,46    train_dataset=train_dataset,47    eval_dataset=eval_dataset,48)49 50trainer.train()51model.save_pretrained("./summarizer_model")52tokenizer.save_pretrained("./summarizer_model")53 54print("Model training complete. Saved to ./summarizer_model")55