removed compute_metrics from training to improve efficiency; added 1000 more dataset; improved inference output quality

Author: codinglabsong

data/subreddit_size_map.json

@@ -4,5 +4,7 @@
     "AskSocialScience": 500,
     "AskDocs": 500,
     "askscience": 1500,
-    "AskHistorians": 1500
+    "AskHistorians": 1500,
+    "AskBiology": 414,
+    "AskEconomics": 532
 }

src/korea_travel_guide/evaluation.py

@@ -3,11 +3,11 @@
 from transformers import EvalPrediction
 
 
-def build_compute_metrics(tok):
+def build_compute_metrics(tok, num_process_workers: int = 2):
     """Return a closure that Hugging Face's Trainer can call."""
     rouge = evaluate.load("rouge")  # longest-substring overlap
     bleu = evaluate.load("bleu")  # n-gram precision
-    # bertscore = evaluate.load("bertscore")  # semantic similarity
+    bertscore = evaluate.load("bertscore")  # semantic similarity
 
     def _compute_metrics(eval_pred: EvalPrediction):
         preds, labels = eval_pred.predictions, eval_pred.label_ids
@@ -23,24 +23,31 @@ def _compute_metrics(eval_pred: EvalPrediction):
 
         # metrics
         rouge_l = rouge.compute(
-            predictions=decoded_preds, references=decoded_labels, use_stemmer=True
+            predictions=decoded_preds,
+            references=decoded_labels,
+            use_stemmer=True,
+            num_process_workers=num_process_workers,
         )["rougeL"]
         bleu_score = bleu.compute(
             predictions=decoded_preds,
             references=[[ref] for ref in decoded_labels],  # BLEU expects list-of-lists
             smooth=True,
+            num_process_workers=num_process_workers,
         )["bleu"]
-        # bert_f1 = np.mean(
-        #     bertscore.compute(
-        #         predictions=decoded_preds, references=decoded_labels, lang="en"
-        #     )["f1"]
-        # )
+        bert_f1 = np.mean(
+            bertscore.compute(
+                predictions=decoded_preds,
+                references=decoded_labels,
+                lang="en",
+                num_process_workers=num_process_workers,
+            )["f1"]
+        )
 
         # round for nice logging
         return {
             "rougeL": round(rouge_l * 100, 4),
             "bleu": round(bleu_score * 100, 4),
-            # "bertscore_f1": round(bert_f1, 4),
+            "bertscore_f1": round(bert_f1, 4),
         }
 
     return _compute_metrics

src/korea_travel_guide/inference.py

@@ -1,6 +1,5 @@
 import logging
 import torch
-from contextlib import nullcontext
 from dataclasses import dataclass, field
 from datasets import load_dataset
 from transformers import (
@@ -34,6 +33,10 @@ class InferenceArgs:
         default_factory=list,
         metadata={"help": "One or more input texts for `predict` mode."},
     )
+    num_process_workers: int = field(
+        default=2,
+        metadata={"help": "Number of workers to parallelize n-gram counting."},
+    )
     use_sdpa_attention: bool = field(
         default=True, metadata={"help": "Enable Sdpa for mem-efficient kernel."}
     )
@@ -82,12 +85,13 @@ def main():
                 output_dir="outputs/inference",
                 per_device_eval_batch_size=inf_args.batch_size,
                 predict_with_generate=True,
+                generation_max_length=384,
                 report_to=[],
             ),
             eval_dataset=ds_tok["test"],
             data_collator=data_collator,
             tokenizer=tok,
-            compute_metrics=build_compute_metrics(tok),
+            compute_metrics=build_compute_metrics(tok, inf_args.num_process_workers),
         )
 
         pred_output = trainer.predict(ds_tok["test"])
@@ -106,14 +110,19 @@ def main():
         )
         enc = {k: v.to(device) for k, v in enc.items()}
 
-        # fast batched generate
+        # fast batched generate (with arguments for higher quality generations)
         out = model.generate(
             **enc,
-            max_length=512,
-            num_beams=5,
-            early_stopping=True,
-            length_penalty=1.0,
-            repetition_penalty=1.1,
+            max_length=200,
+            num_beams=5,  # improves quality
+            do_sample=True,  # add stochasticity
+            length_penalty=1.2,  # >1 favors longer answers
+            repetition_penalty=1.3,  # >1 penalizes reuse of the same token
+            no_repeat_ngram_size=3,  # block exact n-gram repeats
+            top_p=0.9,  # nucleus sampling for diversity
+            temperature=0.8,  # nucleus sampling for diversity
+            early_stopping=True,  # stop on EOS to avoid garbage at the end
+            eos_token_id=tok.eos_token_id,
         )
 
         decoded = tok.batch_decode(out, skip_special_tokens=True)

src/korea_travel_guide/train.py

@@ -2,21 +2,20 @@
 import random
 import numpy as np
 import torch
-from contextlib import nullcontext
 from datasets import load_dataset
 from dataclasses import dataclass, field
 from transformers import (
     HfArgumentParser,
     Seq2SeqTrainingArguments,
     DataCollatorForSeq2Seq,
     Seq2SeqTrainer,
+    EarlyStoppingCallback,
 )
 from typing import List
 from pathlib import Path
 from korea_travel_guide.utils import load_environ_vars, print_trainable_parameters
 from korea_travel_guide.model import build_base_model, build_peft_model
 from korea_travel_guide.data import tokenize_and_format
-from korea_travel_guide.evaluation import build_compute_metrics
 from uuid import uuid4
 
 logger = logging.getLogger(__name__)
@@ -38,6 +37,7 @@ class DataArgs:
         default=False, metadata={"help": "If True, ignore CSVs and load SQuAD instead."}
     )
 
+
 # training & LoRA extras — extend HF’s own Seq2SeqTrainingArguments
 @dataclass
 class CustomTrainingArgs(Seq2SeqTrainingArguments):
@@ -46,22 +46,25 @@ class CustomTrainingArgs(Seq2SeqTrainingArguments):
         default="outputs/bart-base-korea-travel-guide-lora",
         metadata={"help": "Prefix folder for all checkpoints/run logs."},
     )
-    eval_strategy: str = "epoch"
-    save_strategy: str = "epoch"
-    logging_steps: int = 50
-    learning_rate: float = 1e-4
-    lr_scheduler_type: str = "linear"
-    warmup_ratio: float = 0.05
     num_train_epochs: int = 6
     per_device_train_batch_size: int = 8
     per_device_eval_batch_size: int = 16
+    learning_rate: float = 7e-5
+    lr_scheduler_type: str = "cosine"
+    warmup_ratio: float = 0.1
     max_grad_norm: float = 0.5
     label_smoothing_factor: float = 0.1
-    # weight_decay: float = 0.01
-    generation_max_length: int = 384
+    weight_decay: float = 0.01
+
+    eval_strategy: str = "epoch"
+    save_strategy: str = "epoch"
+    logging_steps: int = 50
     save_total_limit: int = 2
+    load_best_model_at_end: bool = True
+    metric_for_best_model: str = "eval/loss"
+    greater_is_better: bool = False
+
     fp16: bool = True
-    predict_with_generate: bool = True
     push_to_hub: bool = False
     report_to: str = "wandb"
     run_name: str = field(
@@ -92,10 +95,6 @@ def parse_args() -> tuple[DataArgs, CustomTrainingArgs]:
     if training_args.push_to_hub and not training_args.hf_hub_repo_id:
         parser.error("--hf_hub_repo_id is required when --push_to_hub is set")
 
-    # # isolate each run’s artefacts (good for sweeps)
-    # run_id = os.environ.get("WANDB_RUN_ID", uuid4().hex[:8])
-    # training_args.output_dir = f"{training_args.output_dir}/{run_id}"
-
     # set wandb for logging
     training_args.report_to = "wandb"
 
@@ -130,14 +129,13 @@ def main() -> None:
     # ---------- Data Preprocessing ----------
     # load either CSVs or SQuAD for a quick pipeline sanity check
     if data_args.use_squad:
-        # 1) pull down SQuAD  
+        # 1) pull down SQuAD
         raw = load_dataset("squad")
+
         # 2) map to simple Q/A pairs (first answer only)
         def to_qa(ex):
-            return {
-                "question": ex["question"],
-                "answer": ex["answers"]["text"][0]
-            }
+            return {"question": ex["question"], "answer": ex["answers"]["text"][0]}
+
         ds = raw.map(to_qa, remove_columns=raw["train"].column_names)
     else:
         # load from your processed CSVs
@@ -198,7 +196,7 @@ def to_qa(ex):
         eval_dataset=ds_tok["validation"],
         tokenizer=tok,
         data_collator=data_collator,
-        # compute_metrics=build_compute_metrics(tok),
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
     )
 
     trainer.train()