import os
from deepspec.trainer import Qwen3DSparkTrainer
BASE_TB_DIR = os.path.expanduser("~/tensorboard")
BASE_CKPT_DIR = os.path.expanduser("~/checkpoints")
project_name = "deepspec"
exp_name = "dspark_block15_qwen36_27b_online"
seed = 42

model = dict(
    target_model_name_or_path="/mnt/t5evo/models/Qwen3.6-27B",
    block_size=15,
    num_draft_layers=5,
    # Qwen3.5-0.8B: 24 layers (6 full-attention), hidden 1024
    target_layer_ids=[1, 16, 31, 46, 61],
    # free embedding slot: tokenizer len 248077 < embed 248320 (same family trick as ornith_9b)
    mask_token_id=248200,
    init_draft_from="/mnt/data/dspark-test/head_q36_27b_dflash_init/model.safetensors",
    num_anchors=512,
    markov_rank=256,
    # match z-lab DFlash warm-start geometry (differs from target-derived defaults)
    head_dim=128,
    num_attention_heads=32,
    num_key_value_heads=8,
    intermediate_size=17408,
    markov_head_type="vanilla",
    confidence_head_alpha=1.0,
    confidence_head_with_markov=True,
    loss_decay_gamma=4.0,
    ce_loss_alpha=0.1,
    l1_loss_alpha=0.9,
)

train = dict(
    trainer_cls=Qwen3DSparkTrainer,
    lr=5.0e-4,
    warmup_ratio=0.04,
    weight_decay=0.0,
    precision="bf16",
    local_batch_size=1,
    global_batch_size=128,   # single Pro 6000; 0.8B target fits with room
    num_train_epochs=5,
    max_train_steps=None,
    max_grad_norm=1.0,
    sharding_strategy="no_shard",
    torch_compile=True,
)

logging = dict(
    logging_steps=10,
    checkpointing_steps=200,
)

data = dict(
    online_target=True,
    target_cache_path=None,
    train_data_paths=[
        "/mnt/data/DeepSpec/train_datasets/q27_regen_12k.jsonl",
    ],
    min_loss_tokens=14,
    chat_template="qwen",
    max_length=2048,
    num_workers=4,
)


def finalize_cfg(cfg):
    logging_cfg = dict(cfg["logging"])
    project_name = str(cfg["project_name"])
    exp_name = str(cfg["exp_name"])
    logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name)
    logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name)
    cfg["logging"] = logging_cfg
    return cfg


# --opts overrides applied at save time
train['num_train_epochs'] = 3