Manually pre-tokenize with apply_chat_template, cap seq length to 512 for OOM fix

train.py

@@ -3,6 +3,7 @@ Full fine-tuning script:
   Model: google/gemma-2-2b-it
   Dataset: talkmap/telecom-conversation-corpus
   Converts turn-based telecom dialogues into alternating-role conversational messages for SFT.
+  Manually tokenizes with truncation to control sequence length and avoid OOM.
 """
 import os
 from collections import defaultdict
@@ -20,8 +21,9 @@ MODEL_ID = "google/gemma-2-2b-it"
 DATASET_ID = "talkmap/telecom-conversation-corpus"
 OUTPUT_DIR = "./gemma-2b-it-telecom"
 HUB_MODEL_ID = "ligaments-dev/gemma-2b-it-telecom"
-MAX_CONVERSATIONS = 5000
-MAX_TURNS = 8  # cap conversation length to reduce activation memory
+MAX_CONVERSATIONS = 10000
+MAX_TURNS = 6
+MAX_SEQ_LENGTH = 512
 
 # ------------------------------------------------------------------
 # Trackio monitoring
@@ -53,11 +55,9 @@ for conv_id in conversations:
     conversations[conv_id].sort(key=lambda x: x["date_time"])
 
 # Convert each conversation into messages format with merged consecutive same-role turns
-# Gemma requires: NO system role, user first, alternating user/assistant
 print("Converting to messages format...")
 messages_data = []
 for conv_id, turns in conversations.items():
-    # Cap turns to MAX_TURNS to keep sequences shorter
     turns = turns[:MAX_TURNS]
 
     messages = []
@@ -75,11 +75,9 @@ for conv_id, turns in conversations.items():
     if current_role is not None:
         messages.append({"role": current_role, "content": "\n".join(current_content)})
 
-    # Gemma requires first turn to be user and alternating roles
     if not messages or messages[0]["role"] != "user":
         continue
 
-    # Verify alternating roles
     valid = True
     for i, msg in enumerate(messages):
         expected_role = "user" if i % 2 == 0 else "assistant"
@@ -89,7 +87,6 @@ for conv_id, turns in conversations.items():
     if not valid:
         continue
 
-    # Ensure conversation ends with assistant (complete pair)
     if messages[-1]["role"] != "assistant":
         continue
 
@@ -98,12 +95,7 @@ for conv_id, turns in conversations.items():
     if len(messages_data) >= MAX_CONVERSATIONS:
         break
 
-train_dataset = Dataset.from_list(messages_data)
-print(f"Total conversations: {len(train_dataset)}")
-
-if len(train_dataset) > 0:
-    print("Sample conversation:")
-    print(train_dataset[0])
+print(f"Total conversations: {len(messages_data)}")
 
 # ------------------------------------------------------------------
 # Tokenizer
@@ -114,6 +106,28 @@ if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
+# ------------------------------------------------------------------
+# Pre-tokenize dataset with truncation (avoids SFTTrainer auto-tokenization)
+# ------------------------------------------------------------------
+print("Pre-tokenizing dataset...")
+
+def apply_and_tokenize(example):
+    try:
+        text = tokenizer.apply_chat_template(
+            example["messages"],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+    except Exception:
+        text = ""
+    return {"text": text}
+
+raw_dataset = Dataset.from_list(messages_data)
+raw_dataset = raw_dataset.map(apply_and_tokenize, remove_columns=["messages"])
+raw_dataset = raw_dataset.filter(lambda x: len(x["text"]) > 0)
+
+print(f"Dataset after filtering: {len(raw_dataset)}")
+
 # ------------------------------------------------------------------
 # Model
 # ------------------------------------------------------------------
@@ -137,6 +151,7 @@ args = SFTConfig(
     per_device_train_batch_size=1,
     gradient_accumulation_steps=4,
     learning_rate=2e-5,
+    max_seq_length=MAX_SEQ_LENGTH,
     logging_strategy="steps",
     logging_steps=10,
     logging_first_step=True,
@@ -155,7 +170,7 @@ print("Initializing SFTTrainer...")
 trainer = SFTTrainer(
     model=model,
     args=args,
-    train_dataset=train_dataset,
+    train_dataset=raw_dataset,
     processing_class=tokenizer,
 )