Create train.py

train.py

@@ -0,0 +1,75 @@
+from datasets import load_dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+    DataCollatorForSeq2Seq,
+)
+import torch
+
+# 1. Load dataset
+dataset = load_dataset("rohitsaxena/MovieSum")
+
+# Rename columns if needed
+dataset = dataset.rename_columns({"script": "input_text", "summary": "target_text"})
+
+# 2. Load model and tokenizer
+model_checkpoint = "facebook/bart-large-cnn"
+tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
+
+# 3. Preprocessing
+def preprocess_function(examples):
+    inputs = tokenizer(
+        examples["input_text"],
+        max_length=1024,
+        padding="max_length",
+        truncation=True,
+    )
+    with tokenizer.as_target_tokenizer():
+        labels = tokenizer(
+            examples["target_text"],
+            max_length=128,
+            padding="max_length",
+            truncation=True,
+        )
+    inputs["labels"] = labels["input_ids"]
+    return inputs
+
+tokenized_dataset = dataset.map(preprocess_function, batched=True)
+
+# 4. Training arguments
+training_args = Seq2SeqTrainingArguments(
+    output_dir="./film-script-summarizer",
+    evaluation_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=2,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    save_total_limit=2,
+    push_to_hub=True,
+    hub_model_id="BhavyaSamhithaMallineni/FilmScriptSummarizer",
+    hub_strategy="every_save",
+    logging_dir="./logs",
+    logging_steps=50,
+    fp16=torch.cuda.is_available(),
+)
+
+# 5. Trainer setup
+data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+trainer = Seq2SeqTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset["train"],
+    eval_dataset=tokenized_dataset["test"],
+    tokenizer=tokenizer,
+    data_collator=data_collator,
+)
+
+# 6. Train and push to hub
+trainer.train()
+trainer.push_to_hub()
+tokenizer.push_to_hub("BhavyaSamhithaMallineni/FilmScriptSummarizer")