Cap turns to 8, require assistant end, 5k conversations for OOM fix

train.py

@@ -3,7 +3,6 @@ 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.
-  NOTE: Gemma chat template does NOT support system role.
 """
 import os
 from collections import defaultdict
@@ -21,7 +20,8 @@ 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 = 10000  # cap for faster training
+MAX_CONVERSATIONS = 5000
+MAX_TURNS = 8  # cap conversation length to reduce activation memory
 
 # ------------------------------------------------------------------
 # Trackio monitoring
@@ -57,6 +57,9 @@ for conv_id in conversations:
 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 = []
     current_role = None
     current_content = []
@@ -86,6 +89,10 @@ for conv_id, turns in conversations.items():
     if not valid:
         continue
 
+    # Ensure conversation ends with assistant (complete pair)
+    if messages[-1]["role"] != "assistant":
+        continue
+
     messages_data.append({"messages": messages})
 
     if len(messages_data) >= MAX_CONVERSATIONS: