Upload train_model.py

train_model.py

@@ -0,0 +1,67 @@
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import TrainingArguments, Trainer
+import os
+import torch
+
+# Load dataset
+ds = load_dataset("knkarthick/dialogsum")
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+
+# Preprocessing function
+def preprocess_function(batch):
+    source = batch['dialogue']
+    target = batch['summary']
+    source_enc = tokenizer(source, padding='max_length', truncation=True, max_length=128)
+    target_enc = tokenizer(target, padding='max_length', truncation=True, max_length=128)
+    labels = target_enc['input_ids']
+    labels = [[(token if token != tokenizer.pad_token_id else -100) for token in label] for label in labels]
+    return {
+        'input_ids': source_enc['input_ids'],
+        'attention_mask': source_enc['attention_mask'],
+        'labels': labels
+    }
+
+# Apply preprocessing
+df_source = ds.map(preprocess_function, batched=True)
+
+# Training arguments
+training_args = TrainingArguments(
+    output_dir='/content/TextSummarizer_output',
+    per_device_train_batch_size=8,
+    num_train_epochs=2,
+    save_total_limit=1,
+    save_strategy="epoch",
+    remove_unused_columns=True,
+    logging_dir='/content/logs',
+    logging_steps=50,
+)
+
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=df_source['train'],
+    eval_dataset=df_source['test'],
+)
+
+# Train
+trainer.train()
+
+# Evaluate
+eval_results = trainer.evaluate()
+print("Evaluation Results:", eval_results)
+
+# ===> Save to Google Drive path
+save_path = "/content/drive/MyDrive/TextSummarizer2/model_directory"
+os.makedirs(save_path, exist_ok=True)
+
+# Save model and tokenizer (use safe_serialization for large model.safetensors)
+model.save_pretrained(save_path, safe_serialization=True)
+tokenizer.save_pretrained(save_path)
+
+print(f"✅ Model and tokenizer saved to: {save_path}")
+print("📦 Files saved:", os.listdir(save_path))