#!/usr/bin/env python3
"""铸渊全自动蒸馏流水线
代码模型完成 → 全自动蒸馏两个1.5B → 深度微调 → 上传COS

使用方法：
  export ZY_OSS_KEY=... ZY_OSS_SECRET=...
  nohup python3 -u scripts/auto_distill_pipeline.py > distill_pipeline.log 2>&1 &
"""
import os, json, time, sys, subprocess, glob
sys.stdout.reconfigure(line_buffering=True)

COS_BUCKET = "sy-finetune-corpus-1317346199"
COS_REGION = "ap-guangzhou"
COS_ID = os.environ.get("ZY_OSS_KEY")
COS_KEY = os.environ.get("ZY_OSS_SECRET")
if not COS_ID or not COS_KEY:
    print("❌ 需要环境变量：export ZY_OSS_KEY=... ZY_OSS_SECRET=...")
    sys.exit(1)

WORK = "/root/autodl-tmp"
TEACHER_MOTHER = f"{WORK}/output/qwen25-7b-sft/final"
TEACHER_CODER = f"{WORK}/output/qwen25-coder-7b-sft/final"
STUDENT = f"{WORK}/cache/Qwen/Qwen2___5-1___5B"
STUDENT_CODER = f"{WORK}/cache/Qwen/Qwen2___5-Coder-1___5B"

def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}"); sys.stdout.flush()
def run(cmd): r = subprocess.run(cmd, shell=True, capture_output=True, text=True); return r.returncode == 0

def cos_check(prefix):
    s = f'from qcloud_cos import CosConfig,CosS3Client; c=CosS3Client(CosConfig(Region="{COS_REGION}",SecretId="{COS_ID}",SecretKey="{COS_KEY}")); r=c.list_objects(Bucket="{COS_BUCKET}",Prefix="{prefix}",MaxKeys=1); print("OK" if "Contents" in r else "NO")'
    return subprocess.run(f'python3 -c "{s}"', shell=True,capture_output=True,text=True).stdout.strip()=="OK"

def cos_dl(key, local):
    s = f'from qcloud_cos import CosConfig,CosS3Client;c=CosS3Client(CosConfig(Region="{COS_REGION}",SecretId="{COS_ID}",SecretKey="{COS_KEY}"));r=c.get_object(Bucket="{COS_BUCKET}",Key="{key}");open("{local}","wb").write(r["Body"].get_raw_stream().read())'
    return run(f'python3 -c "{s}"')

def cos_upload(local, prefix):
    if not os.path.isdir(local): return False
    s = f'import os,glob;from qcloud_cos import CosConfig,CosS3Client;c=CosS3Client(CosConfig(Region="{COS_REGION}",SecretId="{COS_ID}",SecretKey="{COS_KEY}"));count=0\nfor f in glob.glob("{local}/**/*",recursive=True):\n if os.path.isfile(f):\n  rel=os.path.relpath(f,"{local}")\n  c.upload_file(Bucket="{COS_BUCKET}",Key="{prefix}/"+rel,LocalFilePath=f)\n  count+=1\nprint(f"OK:{count}")'
    return run(f'python3 -c "{s}"')

def check_coder():
    log("检查代码模型...")
    if os.path.isdir(TEACHER_CODER): log("  ✅ 本地已存在"); return True
    if cos_check("models/qwen25-coder-7b-sft/final/"): log("  ✅ COS上存在，需手动下载到GPU"); return False
    log("  ⏳ 代码模型尚未完成"); return False

def download_students():
    for name, path in [("Qwen2.5-1.5B",STUDENT),("Qwen2.5-Coder-1.5B",STUDENT_CODER)]:
        if not os.path.isdir(path):
            log(f"下载{name}...")
            run(f"pip3 install modelscope -q && python3 -c \"from modelscope import snapshot_download; snapshot_download('{name.replace('-1.5B','-1___5B').replace('.','/')}', cache_dir='{WORK}/cache')\"")

def run_distill(script, logfile, desc):
    if os.path.isdir(f"{WORK}/output/{desc}/final"): log(f"  ✅ {desc}已存在，跳过"); return True
    log(f"启动{desc}蒸馏...")
    p = f"{WORK}/_run_{desc}.py"
    with open(p,'w') as f: f.write(open(script).read())
    return run(f"cd {WORK} && python3 -u {p} > {logfile} 2>&1")

def run_deepsft(student_path, data, output, logfile, desc):
    if os.path.isdir(f"{output}/final"): log(f"  ✅ {desc}已存在，跳过"); return True
    log(f"启动{desc}深度微调...")
    s = f'''#!/usr/bin/env python3
import os,json,torch,sys;os.environ["CUDA_VISIBLE_DEVICES"]="0"
from transformers import AutoModelForCausalLM,AutoTokenizer,TrainingArguments,Trainer
from datasets import Dataset;from tqdm import tqdm
M="{student_path}";D="{data}";O="{output}"
E,B,G,L,ML=3,4,8,5e-6,2048
os.makedirs(O,exist_ok=True)
raw=[json.loads(l) for l in open(D)]
tok=AutoTokenizer.from_pretrained(M,trust_remote_code=True);tok.pad_token=tok.eos_token
model=AutoModelForCausalLM.from_pretrained(M,trust_remote_code=True,torch_dtype=torch.bfloat16,attn_implementation="sdpa").cuda()
model.config.use_cache=False;model.gradient_checkpointing_enable()
p=[]
for d in tqdm(raw):
 i,l=[],[]
 for m in d["messages"]:
  c=m["content"]
  if not c.strip(): continue
  t=f"<|im_start|>{m['role']}\\n{c}<|im_end|>\\n"
  tk=tok.encode(t,add_special_tokens=False)
  i.extend(tk);l.extend(tk if m["role"]=="assistant" else [-100]*len(tk))
 if len(i)>ML:i,l=i[:ML],l[:ML]
 p.append({{"input_ids":i,"labels":l,"attention_mask":[1]*len(i)}})
def collate(f):
 ml=max(len(x["input_ids"]) for x in f);b={{}}
 for k in["input_ids","labels","attention_mask"]:
  pad=tok.pad_token_id if k!="labels" else -100
  b[k]=torch.tensor([x[k]+[pad]*(ml-len(x[k])) for x in f])
 return b
args=TrainingArguments(output_dir=O,num_train_epochs=E,per_device_train_batch_size=B,gradient_accumulation_steps=G,learning_rate=L,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,gradient_checkpointing=True,optim="adamw_torch",report_to="none")
Trainer(model=model,args=args,train_dataset=Dataset.from_list(p),data_collator=collate).train()
f=os.path.join(O,"final");model.save_pretrained(f);tok.save_pretrained(f)
print("DONE!")
'''
    with open(f"{WORK}/_run_{desc}_sft.py",'w') as f: f.write(s)
    return run(f"cd {WORK} && python3 -u _run_{desc}_sft.py > {logfile} 2>&1")

def main():
    log("="*60); log("铸渊全自动蒸馏流水线 v1.0"); log("="*60)
    
    if not check_coder(): log("代码模型未就绪，终止"); return
    download_students()
    
    # Phase 1: 霜砚蒸馏
    s = os.path.join(os.path.dirname(os.path.abspath(__file__)),"distill_mother.py")
    if os.path.exists(s): run_distill(s,"distill_mother.log","shuangyan-distill")
    else: log("⚠️ distill_mother.py 不在本地，请从仓库获取")
    
    # Phase 2: 铸渊蒸馏 (下载语料+执行)
    log("下载铸渊语料...")
    os.makedirs(f"{WORK}/corpus",exist_ok=True)
    cos_dl("corpus/zhuyuan_full_corpus.jsonl",f"{WORK}/corpus/zhuyuan_full_corpus.jsonl")
    
    s2 = os.path.join(os.path.dirname(os.path.abspath(__file__)),"distill_coder.py")
    if os.path.exists(s2):
        c = open(s2).read().replace('DATA = "/root/autodl-tmp/corpus/zhuyuan_deep_finetune.jsonl"','DATA = "/root/autodl-tmp/corpus/zhuyuan_full_corpus.jsonl"')
        with open(f"{WORK}/_run_zhuyuan-distill.py",'w') as f: f.write(c)
        run_distill(f"{WORK}/_run_zhuyuan-distill.py","distill_coder.log","zhuyuan-distill")
    else: log("⚠️ distill_coder.py 不在本地")
    
    # Phase 3: 铸渊深度SFT
    cos_dl("corpus/zhuyuan_deep_finetune.jsonl",f"{WORK}/corpus/zhuyuan_deep_finetune.jsonl")
    # 合并语料
    with open(f"{WORK}/corpus/zhuyuan_combined.jsonl",'w') as out:
        for src in ["zhuyuan_full_corpus.jsonl","zhuyuan_deep_finetune.jsonl"]:
            p=f"{WORK}/corpus/{src}"
            if os.path.exists(p):
                for l in open(p): out.write(l)
    run_deepsft(f"{WORK}/output/zhuyuan-15b-distill/final",f"{WORK}/corpus/zhuyuan_combined.jsonl",f"{WORK}/output/zhuyuan-15b-deep-sft","zhuyuan_deep_sft.log","zhuyuan")
    
    # Phase 4: 上传
    for local,cos in [("output/shuangyan-15b-distill/final","models/shuangyan-15b-distill/final"),("output/zhuyuan-15b-distill/final","models/zhuyuan-15b-distill/final"),("output/zhuyuan-15b-deep-sft/final","models/zhuyuan-15b-deep-sft/final")]:
        p = f"{WORK}/{local}"
        if os.path.isdir(p):
            log(f"上传 {local}...")
            cos_upload(p,cos)
    
    log("="*60); log("🎉 全自动蒸馏流水线完成！"); log("="*60)

if __name__ == "__main__": main()