code_llm_collector.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Code LLM 數據飛輪系統 (Data Flywheel)
=======================================

使用模型時自動收集數據 → 累積到一定量 → 自動觸發訓練 → 模型變更強

三種收集模式：
  1. 互動模式 — 你問模型寫 code，接受/拒絕/修改回答 → 自動產生訓練數據
  2. Git 監控   — 監控你的 Git repo，新 commit 自動變成訓練數據
  3. API 服務   — 部署成 API，每次請求自動記錄

Usage:
    python code_llm_collector.py chat              # 互動模式
    python code_llm_collector.py watch --repo .    # Git 監控
    python code_llm_collector.py status            # 查看狀態
    python code_llm_collector.py train             # 用收集的數據訓練
    python code_llm_collector.py export            # 匯出到 HuggingFace
"""

import argparse, json, os, subprocess, sys, tempfile, time, hashlib, torch
from datetime import datetime
from pathlib import Path

BASE_MODEL = "Qwen/Qwen2.5-Coder-3B"
ADAPTER_PATH = None
HF_USERNAME = "YOUR_HF_USERNAME"
DATA_DIR = "./collected_data"
SFT_FILE = os.path.join(DATA_DIR, "sft_data.jsonl")
DPO_FILE = os.path.join(DATA_DIR, "dpo_data.jsonl")
GRPO_FILE = os.path.join(DATA_DIR, "grpo_data.jsonl")
META_FILE = os.path.join(DATA_DIR, "metadata.json")
AUTO_TRAIN_THRESHOLD = 100

def ensure_data_dir():
    os.makedirs(DATA_DIR, exist_ok=True)
    if not os.path.exists(META_FILE):
        save_metadata({"total_sft":0,"total_dpo":0,"total_grpo":0,"last_train":None,"train_count":0,"created":datetime.now().isoformat()})

def load_metadata():
    if os.path.exists(META_FILE):
        with open(META_FILE) as f: return json.load(f)
    return {}

def save_metadata(meta):
    with open(META_FILE,"w") as f: json.dump(meta,f,indent=2,ensure_ascii=False)

def append_data(filepath, data):
    with open(filepath,"a",encoding="utf-8") as f: f.write(json.dumps(data,ensure_ascii=False)+"\n")

def count_lines(filepath):
    if not os.path.exists(filepath): return 0
    with open(filepath) as f: return sum(1 for _ in f)

# ============================================================
#  互動模式
# ============================================================
def run_chat():
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    from peft import PeftModel
    ensure_data_dir()
    print("""
    ╔════════════════════════════════════════════════════════════╗
    ║  Code LLM 互動模式 — 邊用邊收集數據                        ║
    ╠════════════════════════════════════════════════════════════╣
    ║  直接輸入問題 → 模型寫 code                                ║
    ║  /accept  → 接受（存為 SFT 數據）                          ║
    ║  /edit    → 貼上修改版（產生 SFT + DPO 對）                 ║
    ║  /reject  → 拒絕                                          ║
    ║  /test    → 加測試（產生 GRPO 數據）                        ║
    ║  /status  → 查看收集狀態                                   ║
    ║  /train   → 用收集的數據訓練                                ║
    ║  /quit    → 退出                                          ║
    ╚════════════════════════════════════════════════════════════╝
    """)
    print("📥 載入模型...")
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
    bnb_config = BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16,bnb_4bit_use_double_quant=True)
    model = AutoModelForCausalLM.from_pretrained(BASE_MODEL,quantization_config=bnb_config,device_map="auto",trust_remote_code=True)
    if ADAPTER_PATH and os.path.exists(ADAPTER_PATH):
        model = PeftModel.from_pretrained(model, ADAPTER_PATH); print(f"   LoRA: {ADAPTER_PATH}")
    model.eval(); print("✅ 模型載入完成\n")
    meta = load_metadata(); current_prompt = None; current_response = None

    while True:
        try: user_input = input("🧑 你: ").strip()
        except (EOFError, KeyboardInterrupt): break
        if not user_input: continue
        if user_input == "/quit": break
        elif user_input == "/status": show_status(); continue
        elif user_input == "/train": trigger_training(); continue
        elif user_input == "/accept":
            if current_prompt and current_response:
                append_data(SFT_FILE, {"messages":[{"role":"user","content":current_prompt},{"role":"assistant","content":current_response}],"timestamp":datetime.now().isoformat(),"source":"chat_accepted"})
                meta["total_sft"] = meta.get("total_sft",0)+1; save_metadata(meta)
                print(f"   ✅ SFT +1 (累計: {meta['total_sft']})"); check_auto_train(meta)
            continue
        elif user_input == "/reject":
            print("   ❌ 已拒絕"); current_response = None; continue
        elif user_input == "/edit":
            if current_prompt and current_response:
                print("   貼上修改後的 code（輸入 END 結束）:")
                edited_lines = []
                while True:
                    line = input()
                    if line.strip() == "END": break
                    edited_lines.append(line)
                edited_code = "\n".join(edited_lines)
                if edited_code.strip():
                    append_data(DPO_FILE, {"prompt":[{"role":"user","content":current_prompt}],"chosen":[{"role":"assistant","content":edited_code}],"rejected":[{"role":"assistant","content":current_response}],"timestamp":datetime.now().isoformat(),"source":"chat_edited"})
                    append_data(SFT_FILE, {"messages":[{"role":"user","content":current_prompt},{"role":"assistant","content":edited_code}],"timestamp":datetime.now().isoformat(),"source":"chat_edited_sft"})
                    meta["total_dpo"] = meta.get("total_dpo",0)+1; meta["total_sft"] = meta.get("total_sft",0)+1; save_metadata(meta)
                    print(f"   ✅ DPO +1 / SFT +1 (DPO:{meta['total_dpo']} SFT:{meta['total_sft']})"); check_auto_train(meta)
            continue
        elif user_input == "/test":
            if current_prompt and current_response:
                print("   貼上 pytest 測試（輸入 END 結束）:")
                test_lines = []
                while True:
                    line = input()
                    if line.strip() == "END": break
                    test_lines.append(line)
                test_code = "\n".join(test_lines)
                if test_code.strip():
                    append_data(GRPO_FILE, {"prompt":[{"role":"user","content":current_prompt}],"solution":current_response,"test":test_code,"timestamp":datetime.now().isoformat(),"source":"chat_test"})
                    meta["total_grpo"] = meta.get("total_grpo",0)+1; save_metadata(meta)
                    print(f"   ✅ GRPO +1 (累計: {meta['total_grpo']})")
            continue

        # 生成回答
        current_prompt = user_input
        messages = [{"role":"system","content":"You are an exceptionally skilled programmer. Write clean, efficient, well-documented code."},{"role":"user","content":user_input}]
        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(text, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7, top_p=0.9, pad_token_id=tokenizer.pad_token_id)
        current_response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        print(f"\n🤖 模型:\n{current_response}\n\n   💡 /accept | /edit | /reject | /test\n")

# ============================================================
#  Git 監控模式
# ============================================================
def run_watch(repo_path):
    ensure_data_dir(); repo_path = os.path.abspath(repo_path)
    print(f"   👀 監控: {repo_path}"); meta = load_metadata()
    seen_file = os.path.join(DATA_DIR, "seen_commits.json")
    seen = set(json.load(open(seen_file))) if os.path.exists(seen_file) else set()
    print(f"   已處理: {len(seen)} commits\n   監控中... (Ctrl+C 停止)\n")
    while True:
        try:
            r = subprocess.run(["git","log","--oneline","-20","--format=%H %s"], cwd=repo_path, capture_output=True, text=True)
            for line in r.stdout.strip().split("\n"):
                if not line.strip(): continue
                parts = line.split(" ",1); h = parts[0]; msg = parts[1] if len(parts)>1 else ""
                if h in seen: continue
                dr = subprocess.run(["git","diff",f"{h}~1",h,"--name-only"], cwd=repo_path, capture_output=True, text=True)
                for f in [x for x in dr.stdout.strip().split("\n") if x.endswith(".py")]:
                    try:
                        fr = subprocess.run(["git","show",f"{h}:{f}"], cwd=repo_path, capture_output=True, text=True)
                        code = fr.stdout
                        if 50 < len(code) < 10000:
                            append_data(SFT_FILE, {"messages":[{"role":"user","content":f"Write: {f}\nCommit: {msg}"},{"role":"assistant","content":code}],"timestamp":datetime.now().isoformat(),"source":"git","commit":h[:8],"file":f})
                            meta["total_sft"] = meta.get("total_sft",0)+1
                            print(f"   📝 {h[:8]} | {f} → SFT ({meta['total_sft']})")
                    except: pass
                seen.add(h)
            save_metadata(meta); json.dump(list(seen), open(seen_file,"w")); check_auto_train(meta)
            time.sleep(30)
        except KeyboardInterrupt: print("\n⏹️ 已停止"); break
        except Exception as e: print(f"   ⚠️ {e}"); time.sleep(30)

def show_status():
    ensure_data_dir(); meta = load_metadata()
    s,d,g = count_lines(SFT_FILE), count_lines(DPO_FILE), count_lines(GRPO_FILE); t = s+d+g
    print(f"""
    📊 數據收集狀態
    ─────────────────────────────
    SFT:  {s:>5} 條  {'█'*min(s//5,30)}
    DPO:  {d:>5} 條  {'█'*min(d//5,30)}
    GRPO: {g:>5} 條  {'█'*min(g//5,30)}
    ─────────────────────────────
    總計: {t:>5} 條
    自動訓練門檻: {AUTO_TRAIN_THRESHOLD} 條
    距下次訓練:   {max(0,AUTO_TRAIN_THRESHOLD-t)} 條
    已訓練次數:   {meta.get('train_count',0)} 次
    """)

def check_auto_train(meta):
    total = count_lines(SFT_FILE)+count_lines(DPO_FILE)+count_lines(GRPO_FILE)
    new = total - meta.get("last_train_total",0)
    if new >= AUTO_TRAIN_THRESHOLD:
        print(f"\n   🔔 累積 {new} 條新數據！運行 python code_llm_collector.py train")

def trigger_training():
    ensure_data_dir(); meta = load_metadata()
    s,d = count_lines(SFT_FILE), count_lines(DPO_FILE)
    if s+d == 0: print("   ⚠️ 無數據"); return
    print(f"\n🚀 訓練中... SFT:{s} DPO:{d}")
    from datasets import Dataset
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model, PeftModel
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
    bnb = BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_compute_dtype=torch.bfloat16,bnb_4bit_use_double_quant=True)
    if s > 0:
        from trl import SFTTrainer, SFTConfig
        data = [json.loads(l) for l in open(SFT_FILE)]; ds = Dataset.from_list([{"messages":x["messages"]} for x in data])
        model = AutoModelForCausalLM.from_pretrained(BASE_MODEL,quantization_config=bnb,device_map="auto",trust_remote_code=True)
        if ADAPTER_PATH and os.path.exists(ADAPTER_PATH): model = PeftModel.from_pretrained(model,ADAPTER_PATH,is_trainable=True)
        else: model = prepare_model_for_kbit_training(model); model = get_peft_model(model, LoraConfig(r=16,lora_alpha=32,lora_dropout=0.05,bias="none",task_type="CAUSAL_LM",target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"]))
        td = os.path.join(DATA_DIR,f"train_{datetime.now().strftime('%Y%m%d_%H%M')}")
        trainer = SFTTrainer(model=model,args=SFTConfig(output_dir=td,learning_rate=2e-4,num_train_epochs=3,per_device_train_batch_size=1,gradient_accumulation_steps=8,max_seq_length=1024,gradient_checkpointing=True,bf16=True,optim="paged_adamw_8bit",logging_steps=10,save_total_limit=1,logging_strategy="steps",logging_first_step=True),processing_class=tokenizer,train_dataset=ds)
        trainer.train(); trainer.save_model(td); print(f"   ✅ SFT: {td}"); del model; torch.cuda.empty_cache()
    meta["last_train"]=datetime.now().isoformat(); meta["train_count"]=meta.get("train_count",0)+1; meta["last_train_total"]=s+d; save_metadata(meta)
    print(f"\n🎉 第 {meta['train_count']} 次訓練完成！")

def export_dataset():
    ensure_data_dir(); s,d = count_lines(SFT_FILE), count_lines(DPO_FILE)
    if s+d == 0: print("   ⚠️ 無數據"); return
    from datasets import Dataset
    if s > 0:
        ds = Dataset.from_list([json.loads(l) for l in open(SFT_FILE)]); n = f"{HF_USERNAME}/my-code-sft-data"
        ds.push_to_hub(n, private=True); print(f"   ✅ SFT: https://huggingface.co/datasets/{n}")
    if d > 0:
        ds = Dataset.from_list([json.loads(l) for l in open(DPO_FILE)]); n = f"{HF_USERNAME}/my-code-dpo-data"
        ds.push_to_hub(n, private=True); print(f"   ✅ DPO: https://huggingface.co/datasets/{n}")

def main():
    parser = argparse.ArgumentParser(description="Code LLM 數據飛輪")
    parser.add_argument("mode", choices=["chat","watch","status","train","export"])
    parser.add_argument("--repo", type=str, default=".")
    args = parser.parse_args()
    {"chat":run_chat,"watch":lambda:run_watch(args.repo),"status":show_status,"train":trigger_training,"export":export_dataset}[args.mode]()

if __name__ == "__main__": main()