#!/usr/bin/env python3
"""
小霜砚LoRA微调 v2 — 用sft_v2.jsonl (1868条霜砚对话)
基座：1.5B蒸馏模板 (Track1)
"""
import os, json, torch
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'

BASE = '/root/autodl-tmp/output/qwen25-15b-shuangyan-distill/final'
OUT = '/root/autodl-tmp/output/shuangyan-v2'
DATA = '/root/autodl-tmp/data/sft_v2.jsonl'
os.makedirs(OUT, exist_ok=True)

print('[1/4] 加载语料...', flush=True)
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model

with open(DATA) as f:
    raw = [json.loads(l) for l in f]
# 去掉system message (cc-002)
data = []
for d in raw:
    msgs = [m for m in d['messages'] if m['role'] != 'system']
    if len(msgs) >= 2:
        data.append({'messages': msgs})
print(f'  加载 {len(data)} 条对话', flush=True)

print('[2/4] 加载模型...', flush=True)
tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    BASE, trust_remote_code=True, torch_dtype=torch.bfloat16,
    attn_implementation='sdpa').cuda()
model.config.use_cache = False

lora = LoraConfig(r=16, lora_alpha=32,
    target_modules=['q_proj','k_proj','v_proj','o_proj'],
    lora_dropout=0.05, bias='none', task_type='CAUSAL_LM')
model = get_peft_model(model, lora)
model.print_trainable_parameters()

print('[3/4] Tokenize...', flush=True)
MAX_LEN = 2048
processed = []
for d in data:
    ii, ll = [], []
    for m in d['messages']:
        c = m['content']
        if not c.strip(): continue
        txt = '<|im_start|>' + m['role'] + '\n' + c + '<|im_end|>\n'
        tk = tokenizer.encode(txt, add_special_tokens=False)
        ii.extend(tk)
        ll.extend(tk if m['role'] == 'assistant' else [-100]*len(tk))
    if len(ii) > MAX_LEN:
        ii, ll = ii[:MAX_LEN], ll[:MAX_LEN]
    if len(ii) > 10:
        processed.append({'input_ids': ii, 'labels': ll, 'attention_mask': [1]*len(ii)})

print(f'  {len(processed)} 条训练数据', flush=True)
ds = Dataset.from_list(processed)

def collate(features):
    ml = 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]*(ml-len(f[k])) for f in features])
    return batch

args = TrainingArguments(
    output_dir=OUT, num_train_epochs=5,
    per_device_train_batch_size=2, gradient_accumulation_steps=4,
    learning_rate=2e-4, warmup_ratio=0.05, lr_scheduler_type='cosine',
    bf16=True, logging_steps=10, save_strategy='epoch', save_total_limit=2,
    remove_unused_columns=False, report_to='none',
    gradient_checkpointing=True, optim='adamw_torch',
)

print('[4/4] 开始微调！', flush=True)
trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=collate)
trainer.train()

fnl = os.path.join(OUT, 'final')
model.save_pretrained(fnl)
tokenizer.save_pretrained(fnl)
print(f'✅ 小霜砚微调完成！{fnl}', flush=True)