#!/usr/bin/env python3
"""光湖母模型→1.5B霜砚模板 蒸馏脚本
Teacher: Qwen2.5-7B (SFT后的母模型)
Student: Qwen2.5-1.5B (将学会霜砚的思维方式)

蒸馏方法：软蒸馏 (KL散度) + 混合SFT

使用方法：
  nohup python3 -u distill_mother.py > distill_mother.log 2>&1 &

配置：
  - 模型路径需根据实际存储位置修改
  - Teacher路径：本地或COS上的SFT输出
  - Student路径：ModelScope/HuggingFace原始模型
"""

import os, json, torch, sys
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from datasets import Dataset
from tqdm import tqdm
import torch.nn.functional as F

# ========== 配置 ==========
TEACHER_PATH = "/root/autodl-tmp/output/qwen25-7b-sft/final"  # 母模型SFT输出
STUDENT_PATH = "/root/autodl-tmp/cache/Qwen/Qwen2___5-1___5B"  # 1.5B学生
DATA = "/root/autodl-tmp/data/sft.jsonl"  # 主语料（也可用shuangyan专属语料）
OUT = "/root/autodl-tmp/output/qwen25-15b-shuangyan-distill"

EPOCHS = 3
BS = 4  # 1.5B可以更大batch
GA = 8
LR = 1e-5
MAX_LEN = 2048
TEMP = 2.0  # 蒸馏温度（越高分布越平滑）
ALPHA = 0.7  # 蒸馏loss权重 (0.7蒸馏 + 0.3SFT)

os.makedirs(OUT, exist_ok=True)

# ========== 1. 加载数据 ==========
print("[1/6] Loading data...")
with open(DATA) as f:
    raw = [json.loads(line) for line in f]
raw = [{"messages": [m for m in obj["messages"] if m["role"] != "system"]} for obj in raw]
print(f"  {len(raw)} examples")

# ========== 2. 加载Teacher + Student ==========
print("[2/6] Loading teacher (7B) and student (1.5B)...")

tokenizer = AutoTokenizer.from_pretrained(STUDENT_PATH, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

print("  Loading teacher...")
teacher = AutoModelForCausalLM.from_pretrained(
    TEACHER_PATH, trust_remote_code=True,
    torch_dtype=torch.bfloat16, attn_implementation="sdpa",
).cuda()
teacher.eval()
for p in teacher.parameters():
    p.requires_grad = False  # Teacher不训练
print(f"  Teacher: {sum(p.numel() for p in teacher.parameters())/1e9:.2f}B")

print("  Loading student...")
student = AutoModelForCausalLM.from_pretrained(
    STUDENT_PATH, trust_remote_code=True,
    torch_dtype=torch.bfloat16, attn_implementation="sdpa",
).cuda()
student.train()
print(f"  Student: {sum(p.numel() for p in student.parameters())/1e9:.2f}B")

# ========== 3. Tokenize（生成teacher logits） ==========
print("[3/6] Tokenizing data and generating teacher logits...")

processed = []
for d in tqdm(raw, desc="Tokenize+Teacher"):
    ids, labs = [], []
    for msg in d["messages"]:
        c = msg["content"]
        if not c.strip(): continue
        t = f"<|im_start|>{msg['role']}\n{c}<|im_end|>\n"
        tok = tokenizer.encode(t, add_special_tokens=False)
        ids.extend(tok)
        labs.extend(tok if msg["role"] == "assistant" else [-100] * len(tok))
    if len(ids) > MAX_LEN:
        ids, labs = ids[:MAX_LEN], labs[:MAX_LEN]
    
    # Teacher生成logits（蒸馏目标）
    with torch.no_grad():
        inp = torch.tensor([ids]).cuda()
        t_out = teacher(input_ids=inp)
        t_logits = t_out.logits[0].float().cpu()  # [seq_len, vocab_size]
    
    processed.append({
        "input_ids": ids, "labels": labs, "attention_mask": [1]*len(ids),
        "teacher_logits": t_logits  # 保存teacher的logits
    })

ds = Dataset.from_list(processed)
total_tok = sum(len(d["input_ids"]) for d in processed)
print(f"  Dataset: {len(ds)} ex, {total_tok:,} tokens")

# ========== 4. 配置训练 ==========
print("[4/6] Training config...")

def distill_collate(features):
    """collate函数：处理padding + 蒸馏loss计算"""
    max_len = max(len(f["input_ids"]) for f in features)
    batch = {}
    for k in ["input_ids", "labels", "attention_mask"]:
        pad = tokenizer.pad_token_id if k != "labels" else -100
        batch[k] = torch.tensor([f[k] + [pad]*(max_len-len(f[k])) for f in features])
    # teacher_logits需要特殊padding（用0填充）
    vocab_size = features[0]["teacher_logits"].size(-1)
    tl = []
    for f in features:
        t = f["teacher_logits"]
        pad_len = max_len - t.size(0)
        if pad_len > 0:
            tl.append(torch.cat([t, torch.zeros(pad_len, vocab_size)], dim=0))
        else:
            tl.append(t[:max_len])
    batch["teacher_logits"] = torch.stack(tl)
    return batch

class DistillTrainer(Trainer):
    """自定义Trainer：蒸馏loss + SFT loss混合"""
    
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        # 前向传播
        outputs = model(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            use_cache=False,
        )
        student_logits = outputs.logits  # [batch, seq_len, vocab_size]
        
        # SFT loss (交叉熵，只计算assistant部分)
        shift_logits = student_logits[..., :-1, :].contiguous()
        shift_labels = inputs["labels"][..., 1:].contiguous()
        sft_loss = F.cross_entropy(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
            ignore_index=-100,
            reduction="mean",
        )
        
        # KL蒸馏loss（teacher vs student）
        teacher_logits = inputs["teacher_logits"]  # [batch, seq_len, vocab_size]
        # 只对assistant部分计算KL
        mask = (inputs["labels"] != -100).unsqueeze(-1).float()  # [batch, seq_len, 1]
        
        # 软化分布
        s_logits_soft = student_logits / TEMP
        t_logits_soft = teacher_logits / TEMP
        
        kl_loss = F.kl_div(
            F.log_softmax(s_logits_soft, dim=-1),
            F.softmax(t_logits_soft, dim=-1),
            reduction="none",
        )
        kl_loss = (kl_loss * mask).sum() / mask.sum()
        kl_loss = kl_loss * (TEMP ** 2)  # 温度缩放
        
        # 混合loss
        total_loss = ALPHA * kl_loss + (1 - ALPHA) * sft_loss
        
        return total_loss

args = TrainingArguments(
    output_dir=OUT, num_train_epochs=EPOCHS,
    per_device_train_batch_size=BS, gradient_accumulation_steps=GA,
    learning_rate=LR, warmup_ratio=0.05, lr_scheduler_type="cosine",
    bf16=True, tf32=True, logging_steps=10,
    save_strategy="epoch", save_total_limit=3,
    remove_unused_columns=False, dataloader_num_workers=4,
    gradient_checkpointing=True, optim="adamw_torch",
    report_to="none", ddp_find_unused_parameters=False,
)

trainer = DistillTrainer(
    model=student, args=args,
    train_dataset=ds, data_collator=distill_collate,
)

# ========== 5. 启动训练 ==========
print("[5/6] Starting distillation!")
gpu = torch.cuda.get_device_name(0)
mem = torch.cuda.get_device_properties(0).total_memory / 1e9
t_params = sum(p.numel() for p in teacher.parameters())
s_params = sum(p.numel() for p in student.parameters())
print(f"  GPU: {gpu} ({mem:.1f}GB)")
print(f"  Teacher: {t_params/1e9:.2f}B | Student: {s_params/1e9:.2f}B")
print(f"  Temp={TEMP}, Alpha={ALPHA}, Eff batch={BS*GA}, LR={LR}")
sys.stdout.flush()

trainer.train()

# ========== 6. 保存 ==========
print("[6/6] Saving distilled model...")
final = os.path.join(OUT, "final")
trainer.save_model(final)
tokenizer.save_pretrained(final)

# ⚠️ 关键修复：Qwen chat template 使用 <|im_end|> (151645) 作为对话EOS
# 但默认 eos_token_id=151643 (<|endoftext|>)
# 不修复会导致部署时模型无限生成 → 死循环乱码
# 注意：必须同时修复 config.json 和 generation_config.json！
model.config.eos_token_id = 151645
model.config.save_pretrained(final)

model.generation_config.eos_token_id = 151645
model.generation_config.pad_token_id = 151645
model.generation_config.save_pretrained(final)

# 修复 tokenizer 默认system prompt（避免 "You are Qwen..."）
import json as _json
_tok_cfg_path = os.path.join(final, "tokenizer_config.json")
with open(_tok_cfg_path) as _f:
    _tok_cfg = _json.load(_f)
_tok_cfg["default_system"] = ""
with open(_tok_cfg_path, "w") as _f:
    _json.dump(_tok_cfg, _f, indent=2, ensure_ascii=False)

peak = torch.cuda.max_memory_allocated() / 1e9
print(f"  Model: {final}")
print(f"  Peak VRAM: {peak:.2f}GB / {mem:.1f}GB")
print("DONE!")