# =============================================================================

# train.py

# Dataset Preparation + Finetuning Entry Point

# SmolLM2 Service Space

# Copyright 2026 - Volkan Kücükbudak

# Apache License V2 + ESOL 1.1

# =============================================================================

# Usage:

# python train.py --mode export → export HF dataset to training format

# python train.py --mode validate → validate ADI weights against dataset

# python train.py --mode finetune → finetune SmolLM2 on exported data

# =============================================================================

import os

import argparse

import json

import logging

from datetime import datetime

from pathlib import Path

# ── Path Resolution ───────────────────────────────────────────────────────────

# HF Spaces: /tmp/ (read-only filesystem)

# Local dev: current directory

_TMP = Path("/tmp") if os.getenv("SPACE_ID") else Path(".")

TRAIN_DATA = _TMP / "train_data.jsonl"

VALID_RESULT = _TMP / "validation_results.json"

MODEL_OUTPUT = _TMP / "finetuned_model"

import model as model_module

from adi import DumpindexAnalyzer

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")

logger = logging.getLogger("train")

# =============================================================================

# Mode 1 — Export dataset to training format

# =============================================================================

def export_dataset(output_path: str = None):

"""

Export HF dataset logs to JSONL format for training.

Includes HIGH_PRIORITY, MEDIUM_PRIORITY and BLOCKED entries.

BLOCKED entries teach the model what to reject.

REJECT entries (ADI noise/quality fail) are skipped — no response logged.

"""

output = Path(output_path) if output_path else TRAIN_DATA

logger.info("Loading dataset from HF...")

entries = model_module.load_logs()

# ── DEBUG: remove after fix ───────────────────────────────────────────

#if entries:

#logger.info(f"Keys: {list(entries[0].keys())}")

#logger.info(f"Sample: {entries[0]}")

# ─────────────────────────────────────────────────────────────────────

if not entries:

logger.warning("Dataset empty — nothing to export")

return

count = 0

skipped = 0

with open(output, "w") as f:

for entry in entries:

# Skip ADI-rejected entries — no meaningful response logged

if entry.get("adi_decision") == "REJECT":

skipped += 1

continue

if not entry.get("response"):

skipped += 1

continue

# Format as instruction tuning pair

# BLOCKED entries are included — model learns what to refuse

record = {

"instruction": entry.get("system_prompt", "You are a helpful assistant."),

"input": entry.get("prompt", ""),

"output": entry.get("response", ""),

"adi_score": entry.get("adi_score"),

"adi_decision": entry.get("adi_decision"),

"is_safe": entry.get("adi_decision") != "BLOCKED",

}

f.write(json.dumps(record) + "\n")

count += 1

logger.info(f"Exported {count}/{len(entries)} entries → {output} (skipped: {skipped})")

# =============================================================================

# Mode 2 — Validate ADI weights against collected data

# =============================================================================

def validate_adi():

"""

Run ADI weight validation against dataset.

Uses entries that have human_label field (manually labeled).

"""

logger.info("Loading dataset for ADI validation...")

entries = model_module.load_logs()

labeled = [(e["prompt"], e["human_label"]) for e in entries if e.get("human_label")]

if not labeled:

logger.warning("No labeled entries found — add 'human_label' field to dataset entries")

logger.info("Expected labels: REJECT | MEDIUM_PRIORITY | HIGH_PRIORITY")

return

analyzer = DumpindexAnalyzer()

accuracy = analyzer.validate_weights(labeled)

logger.info(f"ADI Validation accuracy: {accuracy:.1%} on {len(labeled)} samples")

result = {

"timestamp": datetime.utcnow().isoformat(),

"accuracy": accuracy,

"samples": len(labeled),

"weights": analyzer.weights,

}

VALID_RESULT.write_text(json.dumps(result, indent=2))

logger.info(f"Results saved → {VALID_RESULT}")

# =============================================================================

# Mode 3 — Finetune SmolLM2 with TRL SFTTrainer

# =============================================================================

def finetune():

"""

Finetune SmolLM2 on exported dataset using TRL SFTTrainer.

Requires export first + enough data (500+ samples recommended).

On completion: pushes finetuned weights to private HF model repo.

"""

if not TRAIN_DATA.exists():

logger.error(f"train_data.jsonl not found at {TRAIN_DATA} — run export first")

return

lines = TRAIN_DATA.read_text().strip().splitlines()

logger.info(f"Training samples available: {len(lines)}")

if len(lines) < 10:

logger.error(f"Too few samples ({len(lines)}) — aborting finetune")

return

if len(lines) < 500:

logger.warning(f"Only {len(lines)} samples — recommend 500+ for meaningful finetuning")

# ── Imports ───────────────────────────────────────────────────────────────

try:

from transformers import AutoModelForCausalLM, AutoTokenizer

from trl import SFTTrainer, SFTConfig

from datasets import Dataset

import torch

except ImportError as e:

logger.error(f"Missing dependency: {e} — run: pip install trl transformers datasets torch")

return

# ── Load dataset ──────────────────────────────────────────────────────────

logger.info("Loading training data...")

records = [json.loads(l) for l in lines]

def format_record(record):

"""Format record into chat template string."""

instruction = record.get("instruction", "You are a helpful assistant.")

user_input = record.get("input", "")

output = record.get("output", "")

return {

"text": f"<|system|>\n{instruction}\n<|user|>\n{user_input}\n<|assistant|>\n{output}"

}

formatted = [format_record(r) for r in records]

dataset = Dataset.from_list(formatted)

logger.info(f"Dataset ready: {len(dataset)} samples")

# ── Load model + tokenizer ────────────────────────────────────────────────

model_id = model_module.get_model_id()

kwargs = model_module.get_model_kwargs()

device = "cuda" if torch.cuda.is_available() else "cpu"

logger.info(f"Loading base model: {model_id} on {device}...")

tokenizer = AutoTokenizer.from_pretrained(model_id, **kwargs)

model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)

# Ensure pad token exists

if tokenizer.pad_token is None:

tokenizer.pad_token = tokenizer.eos_token

# ── Training config ───────────────────────────────────────────────────────

# Conservative settings for CPU / low RAM (2-8GB)

sft_config = SFTConfig(

output_dir=str(MODEL_OUTPUT),

num_train_epochs=3,

per_device_train_batch_size=1, # CPU friendly

gradient_accumulation_steps=4, # effective batch size = 4

learning_rate=2e-5,

warmup_steps=10,

logging_steps=10,

save_steps=50,

save_total_limit=2,

fp16=False, # no GPU, no fp16

bf16=False,

dataloader_num_workers=0, # HF Spaces: no multiprocessing

report_to="none", # no wandb/tensorboard

max_seq_length=512, # SmolLM2 context limit

dataset_text_field="text",

)

# ── SFTTrainer ────────────────────────────────────────────────────────────

logger.info("Initializing SFTTrainer...")

trainer = SFTTrainer(

model=model,

args=sft_config,

train_dataset=dataset,

#tokenizer=tokenizer,

)

# ── Train ─────────────────────────────────────────────────────────────────

logger.info("Starting finetuning...")

start = datetime.utcnow()

trainer.train()

duration = (datetime.utcnow() - start).total_seconds()

logger.info(f"Training complete in {duration:.0f}s")

# ── Save locally ──────────────────────────────────────────────────────────

trainer.save_model(str(MODEL_OUTPUT))

tokenizer.save_pretrained(str(MODEL_OUTPUT))

logger.info(f"Model saved → {MODEL_OUTPUT}")

# ── Push to HF private repo ───────────────────────────────────────────────

token = model_module.TOKEN

private_repo = model_module.PRIVATE_MODEL

if token and private_repo:

logger.info(f"Pushing to HF: {private_repo}...")

try:

model.push_to_hub(private_repo, token=token, private=True)

tokenizer.push_to_hub(private_repo, token=token, private=True)

model_module.push_model_card({

"model_id": model_id,

"samples": len(dataset),

"epochs": 3,

"duration_sec": int(duration),

"finetuned_from": model_id,

})

logger.info(f"Model pushed → {private_repo}")

except Exception as e:

logger.error(f"Push failed: {type(e).__name__}: {e}")

else:

logger.warning("No token or private repo configured — skipping HF push")

# =============================================================================

# CLI

# =============================================================================

if __name__ == "__main__":

parser = argparse.ArgumentParser(description="SmolLM2 Training Utilities")

parser.add_argument(

"--mode",

choices=["export", "validate", "finetune"],

required=True,

help="export: dump dataset to JSONL | validate: test ADI weights | finetune: train model"

)

parser.add_argument("--output", default=None, help="Output file for export mode (default: auto)")

args = parser.parse_args()

if args.mode == "export":

export_dataset(args.output)

elif args.mode == "validate":

validate_adi()

elif args.mode == "finetune":

finetune()