Add training script

train.py

@@ -0,0 +1,119 @@
+import os
+import sys
+from datasets import load_dataset, Dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from trl import SFTTrainer, SFTConfig
+import trackio
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+MODEL_ID = "google/gemma-2b-it"
+DATASET_ID = "talkmap/telecom-conversation-corpus"
+OUTPUT_REPO = "ligaments-dev/gemma-2b-telecom-sft"
+MAX_LENGTH = 2048
+
+# ---------------------------------------------------------------------------
+# Logging / tracking
+# ---------------------------------------------------------------------------
+trackio.init(project="gemma-telecom-sft")
+
+# ---------------------------------------------------------------------------
+# Load tokenizer
+# ---------------------------------------------------------------------------
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# ---------------------------------------------------------------------------
+# Preprocess dataset: group rows by conversation_id into messages
+# ---------------------------------------------------------------------------
+print("Loading dataset...")
+raw = load_dataset(DATASET_ID, split="train")
+
+# The dataset has columns: conversation_id, speaker, date_time, text
+# We group by conversation_id and build a single messages list per conversation.
+print("Grouping conversations...")
+conv_map = {}
+for ex in raw:
+    cid = ex["conversation_id"]
+    if cid not in conv_map:
+        conv_map[cid] = []
+    role = "user" if ex["speaker"] == "client" else "assistant"
+    conv_map[cid].append({"role": role, "content": ex["text"]})
+
+# Build a huggingface Dataset from the grouped conversations
+messages_list = []
+for cid, msgs in conv_map.items():
+    # Keep only alternating user/assistant; skip system-like turns if any.
+    # Gemma-2b-it chat template expects user/assistant.
+    messages_list.append({"messages": msgs})
+
+train_dataset = Dataset.from_list(messages_list)
+print(f"Prepared {len(train_dataset)} conversation examples.")
+
+# Save processed dataset so we don't recompute on resume
+processed_path = "/tmp/telecom_processed"
+train_dataset.save_to_disk(processed_path)
+print(f"Saved processed dataset to {processed_path}")
+
+# ---------------------------------------------------------------------------
+# Load model
+# ---------------------------------------------------------------------------
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype="auto",
+    device_map="auto",
+)
+
+# ---------------------------------------------------------------------------
+# Training arguments
+# ---------------------------------------------------------------------------
+training_args = SFTConfig(
+    output_dir="/tmp/gemma-telecom-sft",
+    num_train_epochs=3,
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=8,
+    learning_rate=2e-5,
+    warmup_ratio=0.1,
+    lr_scheduler_type="cosine",
+    max_length=MAX_LENGTH,
+    packing=True,
+    gradient_checkpointing=True,
+    bf16=True,
+    logging_steps=10,
+    save_strategy="epoch",
+    push_to_hub=True,
+    hub_model_id=OUTPUT_REPO,
+    hub_private_repo=False,
+    disable_tqdm=True,
+    logging_strategy="steps",
+    logging_first_step=True,
+    report_to=["trackio"],
+)
+
+# ---------------------------------------------------------------------------
+# Trainer
+# ---------------------------------------------------------------------------
+print("Initializing SFTTrainer...")
+trainer = SFTTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    processing_class=tokenizer,
+)
+
+# ---------------------------------------------------------------------------
+# Train
+# ---------------------------------------------------------------------------
+print("Starting training...")
+trainer.train()
+
+# ---------------------------------------------------------------------------
+# Push to hub
+# ---------------------------------------------------------------------------
+print("Pushing model to hub...")
+trainer.push_to_hub(commit_message="Full SFT on telecom conversation corpus")
+print("Done!")