Update train_seallm_khm_sum.py

train_seallm_khm_sum.py

@@ -4,20 +4,22 @@ from datasets import load_dataset
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
-    BitsAndBytesConfig,
     TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+    BitsAndBytesConfig,
 )
-from trl import SFTTrainer
-from peft import LoraConfig
+from peft import LoraConfig, get_peft_model
 
 MODEL_NAME = "SeaLLMs/SeaLLMs-v3-1.5B"
 DATASET_NAME = "bltlab/lr-sum"
 DATASET_CONFIG = "khm"
 
+
 def load_khm_dataset():
     raw = load_dataset(DATASET_NAME, DATASET_CONFIG)
 
-    # Try to find train/validation; if not, split test
+    # Try standard splits first
     if "train" in raw:
         train = raw["train"]
         if "validation" in raw:
@@ -28,7 +30,7 @@ def load_khm_dataset():
             split = train.train_test_split(test_size=0.05, seed=42)
             train, eval_ds = split["train"], split["test"]
     else:
-        # Some LR-Sum subsets only have 'test'; we split that.
+        # Some subsets only have 'test'; split that
         split = raw["test"].train_test_split(test_size=0.1, seed=42)
         train, eval_ds = split["train"], split["test"]
 
@@ -36,7 +38,7 @@ def load_khm_dataset():
         article = example["text"]
         summary = example["summary"]
 
-        # Simple Khmer instruction → Khmer summary
+        # Simple Khmer instruction-style format
         text = (
             "សូមសង្ខេបអត្ថបទខាងក្រោមជាភាសាខ្មែរ៖\n\n"
             f"{article}\n\n"
@@ -62,7 +64,7 @@ def load_khm_dataset():
 
 
 def load_model_and_tokenizer():
-    # QLoRA 4-bit quantization config
+    # QLoRA 4-bit config
     bnb_config = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_use_double_quant=True,
@@ -74,7 +76,6 @@ def load_model_and_tokenizer():
         MODEL_NAME,
         trust_remote_code=True,
     )
-
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
@@ -85,15 +86,16 @@ def load_model_and_tokenizer():
         trust_remote_code=True,
     )
 
-    # Enable gradient checkpointing for memory
     model.gradient_checkpointing_enable()
 
     return model, tokenizer
 
+
 def main():
     train_ds, eval_ds = load_khm_dataset()
     model, tokenizer = load_model_and_tokenizer()
 
+    # Apply LoRA to the model
     lora_config = LoraConfig(
         r=64,
         lora_alpha=16,
@@ -101,8 +103,40 @@ def main():
         bias="none",
         task_type="CAUSAL_LM",
     )
+    model = get_peft_model(model, lora_config)
+
+    # Tokenize datasets
+    max_length = 1024
+
+    def tokenize_function(batch):
+        out = tokenizer(
+            batch["text"],
+            max_length=max_length,
+            truncation=True,
+            padding="max_length",
+        )
+        # Causal LM: labels = input_ids
+        out["labels"] = out["input_ids"].copy()
+        return out
+
+    train_tokenized = train_ds.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["text"],
+        desc="Tokenizing train set",
+    )
+    eval_tokenized = eval_ds.map(
+        tokenize_function,
+        batched=True,
+        remove_columns=["text"],
+        desc="Tokenizing eval set",
+    )
+
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False,
+    )
 
-    # Use standard TrainingArguments instead of SFTConfig
     training_args = TrainingArguments(
         output_dir="seallm-khm-sum-lora",
         num_train_epochs=2,
@@ -115,33 +149,27 @@ def main():
         save_total_limit=2,
         lr_scheduler_type="cosine",
         warmup_ratio=0.03,
-        # old transformers may not support bf16, so let's be safe:
-        fp16=True,          # use fp16 instead of bf16
-        report_to="none",   # if this errors next, we’ll drop it
+        fp16=True,          # safer for old transformers
+        report_to="none",   # remove if this crashes
     )
 
-
-    trainer = SFTTrainer(
+    trainer = Trainer(
         model=model,
-        tokenizer=tokenizer,
-        train_dataset=train_ds,
-        eval_dataset=eval_ds,
-        peft_config=lora_config,
         args=training_args,
-        dataset_text_field="text",
-        max_seq_length=1024,    # set here instead of in config
-        # packing=False,        # keep off for compatibility
+        train_dataset=train_tokenized,
+        eval_dataset=eval_tokenized,
+        data_collator=data_collator,
     )
 
     trainer.train()
 
-    # Save LoRA adapter and tokenizer
-    trainer.model.save_pretrained("seallm-khm-sum-lora")
+    # Save LoRA adapter + tokenizer
+    model.save_pretrained("seallm-khm-sum-lora")
     tokenizer.save_pretrained("seallm-khm-sum-lora")
 
     repo_id = os.environ.get("OUTPUT_REPO_ID", "")
     if repo_id:
-        trainer.model.push_to_hub(repo_id)
+        model.push_to_hub(repo_id)
         tokenizer.push_to_hub(repo_id)