Added more epochs for training the model

model.py

@@ -1,70 +1,70 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from transformers import Trainer, TrainingArguments, GenerationConfig
-
-
-def load_model(model_name="facebook/bart-large-cnn"):
-    """
-    Load a pre-trained summarization model
-    Options: facebook/bart-large-cnn, google/pegasus-xsum, sshleifer/distilbart-cnn-12-6
-    """
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    return model, tokenizer
-
-
-def fine_tune_model(model, tokenizer, dataset, output_dir="./summarization_model"):
-    """Fine-tune model on prepared dataset"""
-    training_args = TrainingArguments(
-        output_dir=output_dir,
-        per_device_train_batch_size=4,
-        per_device_eval_batch_size=4,
-        gradient_accumulation_steps=4,
-        learning_rate=5e-5,
-        num_train_epochs=6,
-        save_strategy="epoch",
-        eval_strategy="epoch",
-        load_best_model_at_end=True,
-        report_to="none",
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=dataset["train"],
-        eval_dataset=dataset["validation"],
-        tokenizer=tokenizer,
-    )
-
-    trainer.train()
-
-    tokenizer.save_pretrained(output_dir)
-    model.save_pretrained(output_dir)
-
-    return model
-
-
-def generate_stylized_summary(text, model, tokenizer, style="formal", max_length=150):
-    """Generate a summary in the specified style"""
-    # Prepend style token to input
-    styled_input = f"[{style.upper()}] {text}"
-    inputs = tokenizer(
-        styled_input, return_tensors="pt", max_length=1024, truncation=True
-    )
-
-    generation_config = GenerationConfig(
-        max_length=max_length,
-        min_length=56,
-        early_stopping=True,
-        num_beams=4,
-        length_penalty=2.0,
-        no_repeat_ngram_size=3,
-        forced_bos_token_id=0,
-    )
-
-    summary_ids = model.generate(
-        inputs["input_ids"], generation_config=generation_config
-    )
-
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    return summary
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import Trainer, TrainingArguments, GenerationConfig
+
+
+def load_model(model_name="facebook/bart-large-cnn"):
+    """
+    Load a pre-trained summarization model
+    Options: facebook/bart-large-cnn, google/pegasus-xsum, sshleifer/distilbart-cnn-12-6
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    return model, tokenizer
+
+
+def fine_tune_model(model, tokenizer, dataset, output_dir="./summarization_model"):
+    """Fine-tune model on prepared dataset"""
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        gradient_accumulation_steps=4,
+        learning_rate=5e-5,
+        num_train_epochs=20,
+        save_strategy="epoch",
+        eval_strategy="epoch",
+        load_best_model_at_end=True,
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["validation"],
+        tokenizer=tokenizer,
+    )
+
+    trainer.train()
+
+    tokenizer.save_pretrained(output_dir)
+    model.save_pretrained(output_dir)
+
+    return model
+
+
+def generate_stylized_summary(text, model, tokenizer, style="formal", max_length=150):
+    """Generate a summary in the specified style"""
+    # Prepend style token to input
+    styled_input = f"[{style.upper()}] {text}"
+    inputs = tokenizer(
+        styled_input, return_tensors="pt", max_length=1024, truncation=True
+    )
+
+    generation_config = GenerationConfig(
+        max_length=max_length,
+        min_length=56,
+        early_stopping=True,
+        num_beams=4,
+        length_penalty=2.0,
+        no_repeat_ngram_size=3,
+        forced_bos_token_id=0,
+    )
+
+    summary_ids = model.generate(
+        inputs["input_ids"], generation_config=generation_config
+    )
+
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return summary