GENERATION_PROMPT = """你是一名{domain}领域的资深专家。请基于以下文档片段，生成 {n} 个高质量的问答对。

要求：
1. 问题必须是真实业务场景中会问的，避免"这段文字讲了什么"这类无意义问题。
2. 答案必须完全基于文档内容，不得引入外部知识。
3. 答案要包含推理过程（如果问题需要推理）。
4. 覆盖不同难度：简单事实抽取、中等归纳、困难推理。
5. 输出 JSON 格式：[{{"question": "...", "answer": "...", "difficulty": "easy|medium|hard", "reasoning": "..."}}]

文档片段：
{chunk}
"""

import re
from collections import Counter
from sentence_transformers import SentenceTransformer
import numpy as np

class DataQualityScorer:
    def __init__(self, embed_model="BAAI/bge-m3"):
        self.embedder = SentenceTransformer(embed_model)

    def score(self, qa_pairs):
        """对问答对列表打分，返回带分数的列表"""
        questions = [p["question"] for p in qa_pairs]
        answers = [p["answer"] for p in qa_pairs]

        # 1. 长度合理性
        len_scores = [
            1.0 if 10 <= len(q) <= 200 and 20 <= len(a) <= 2000 else 0.3
            for q, a in zip(questions, answers)
        ]

        # 2. 语义多样性（问题之间不能太像）
        q_embs = self.embedder.encode(questions, normalize_embeddings=True)
        sim_matrix = q_embs @ q_embs.T
        np.fill_diagonal(sim_matrix, 0)
        max_sim = sim_matrix.max(axis=1)
        diversity_scores = np.clip(1.0 - max_sim, 0, 1).tolist()

        # 3. 答案是否包含具体信息（数字、专有名词）
        info_scores = []
        for a in answers:
            has_num = bool(re.search(r"\d", a))
            has_proper = bool(re.search(r"[A-Z]{2,}|[\u4e00-\u9fa5]{2,}型", a))
            info_scores.append(0.5 + 0.25 * has_num + 0.25 * has_proper)

        # 4. 答案是否拒绝回答（坏样本）
        refusal_patterns = ["无法回答", "不知道", "没有相关信息", "抱歉"]
        refusal_scores = [
            0.1 if any(p in a for p in refusal_patterns) else 1.0
            for a in answers
        ]

        # 加权总分
        final = [
            0.2 * l + 0.3 * d + 0.2 * i + 0.3 * r
            for l, d, i, r in zip(len_scores, diversity_scores, info_scores, refusal_scores)
        ]

        for p, s in zip(qa_pairs, final):
            p["quality_score"] = round(s, 3)
        return qa_pairs

    def filter(self, qa_pairs, threshold=0.65):
        scored = self.score(qa_pairs)
        return [p for p in scored if p["quality_score"] >= threshold]

# train_r1_lora.py
import os
import json
import time
import torch
from dataclasses import dataclass, field
from typing import Optional

from unsloth import FastLanguageModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from transformers import TrainingArguments

# ============ 1. 配置 ============
@dataclass
class TrainConfig:
    model_name: str = "unsloth/DeepSeek-R1-Distill-Qwen-7B"
    max_seq_length: int = 2048
    load_in_4bit: bool = True
    dtype: Optional[torch.dtype] = None

    lora_r: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    target_modules: list = field(default_factory=lambda: [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ])

    dataset_path: str = "./data/train.jsonl"
    val_path: str = "./data/val.jsonl"
    output_dir: str = "./outputs/r1-lora-v1"

    per_device_batch_size: int = 2
    grad_accum: int = 8
    num_epochs: int = 3
    learning_rate: float = 2e-4
    warmup_ratio: float = 0.05
    lr_scheduler: str = "cosine"
    weight_decay: float = 0.01
    max_grad_norm: float = 1.0

    logging_steps: int = 10
    save_steps: int = 200
    eval_steps: int = 100
    save_total_limit: int = 3
    seed: int = 42

cfg = TrainConfig()

# ============ 2. 加载模型 ============
print(f"[INFO] Loading model: {cfg.model_name}")
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=cfg.model_name,
    max_seq_length=cfg.max_seq_length,
    dtype=cfg.dtype,
    load_in_4bit=cfg.load_in_4bit,
    use_gradient_checkpointing="unsloth",
    random_state=cfg.seed,
)

# ============ 3. 注入 LoRA ============
model = FastLanguageModel.get_peft_model(
    model,
    r=cfg.lora_r,
    lora_alpha=cfg.lora_alpha,
    lora_dropout=cfg.lora_dropout,
    target_modules=cfg.target_modules,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=cfg.seed,
    use_rslora=False,
    loftq_config=None,
)

# 打印可训练参数
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"[INFO] Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.3f}%)")

# ============ 4. 数据加载与格式化 ============
tokenizer = get_chat_template(tokenizer, chat_template="deepseek-r1")

def format_sample(examples):
    """把 {question, answer, reasoning} 格式化为 R1 对话格式"""
    texts = []
    for q, a, r in zip(examples["question"], examples["answer"], examples.get("reasoning", [""] * len(examples["question"]))):
        # R1 的思考链格式：<|think|>...<|/think|> 后接答案
        if r:
            assistant_content = f"<|think|>\n{r}\n<|/think|>\n\n{a}"
        else:
            assistant_content = a
        messages = [
            {"role": "user", "content": q},
            {"role": "assistant", "content": assistant_content},
        ]
        texts.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False))
    return {"text": texts}

train_ds = load_dataset("json", data_files=cfg.dataset_path, split="train")
val_ds = load_dataset("json", data_files=cfg.val_path, split="train")
train_ds = train_ds.map(format_sample, batched=True, remove_columns=train_ds.column_names)
val_ds = val_ds.map(format_sample, batched=True, remove_columns=val_ds.column_names)
print(f"[INFO] Train samples: {len(train_ds)}, Val samples: {len(val_ds)}")

# ============ 5. 显存监控回调 ============
from transformers import TrainerCallback

class MemoryCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if torch.cuda.is_available():
            alloc = torch.cuda.memory_allocated() / 1024**3
            reserved = torch.cuda.memory_reserved() / 1024**3
            peak = torch.cuda.max_memory_allocated() / 1024**3
            print(f"[MEM] step={state.global_step} alloc={alloc:.2f}G reserved={reserved:.2f}G peak={peak:.2f}G")

# ============ 6. 训练配置 ============
sft_config = SFTConfig(
    output_dir=cfg.output_dir,
    per_device_train_batch_size=cfg.per_device_batch_size,
    per_device_eval_batch_size=cfg.per_device_batch_size,
    gradient_accumulation_steps=cfg.grad_accum,
    num_train_epochs=cfg.num_epochs,
    learning_rate=cfg.learning_rate,
    warmup_ratio=cfg.warmup_ratio,
    lr_scheduler_type=cfg.lr_scheduler,
    weight_decay=cfg.weight_decay,
    max_grad_norm=cfg.max_grad_norm,
    logging_steps=cfg.logging_steps,
    save_steps=cfg.save_steps,
    eval_steps=cfg.eval_steps,
    eval_strategy="steps",
    save_strategy="steps",
    save_total_limit=cfg.save_total_limit,
    bf16=is_bfloat16_supported(),
    fp16=not is_bfloat16_supported(),
    optim="adamw_8bit",
    seed=cfg.seed,
    report_to="none",
    dataset_text_field="text",
    max_seq_length=cfg.max_seq_length,
    packing=True,  # 序列打包，吞吐翻倍
    dataset_num_proc=4,
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_ds,
    eval_dataset=val_ds,
    args=sft_config,
    callbacks=[MemoryCallback()],
)

# ============ 7. 训练 ============
print("[INFO] Start training...")
t0 = time.time()
train_result = trainer.train(resume_from_checkpoint=None)
elapsed = time.time() - t0
print(f"[INFO] Training done in {elapsed/60:.1f} min")
print(f"[INFO] Final train loss: {train_result.training_loss:.4f}")

# ============ 8. 保存 ============
trainer.save_model(cfg.output_dir)
tokenizer.save_pretrained(cfg.output_dir)
print(f"[INFO] Model saved to {cfg.output_dir}")

# 保存训练指标
metrics = train_result.metrics
metrics["train_time_min"] = elapsed / 60
with open(os.path.join(cfg.output_dir, "train_metrics.json"), "w") as f:
    json.dump(metrics, f, indent=2)

# eval_generation.py
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="./outputs/r1-lora-v1",
    max_seq_length=2048,
    load_in_4bit=True,
)
FastLanguageModel.for_inference(model)
tokenizer = get_chat_template(tokenizer, chat_template="deepseek-r1")

# 三类测试样本
test_cases = {
    "in_domain": [
        "患者男性 65 岁，主诉胸痛 2 小时，心电图 ST 段抬高，如何初步处理？",
    ],
    "out_domain": [
        "用 Python 写一个快速排序",
        "解释一下量子纠缠",
    ],
    "edge_case": [
        "今天天气怎么样？",  # 领域外，应该拒答或通用回答
        "",  # 空输入
        "A" * 500,  # 超长无意义输入
    ],
}

for category, prompts in test_cases.items():
    print(f"\n{'='*60}\n[{category}]\n{'='*60}")
    for p in prompts:
        messages = [{"role": "user", "content": p}]
        inputs = tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True, return_tensors="pt"
        ).to("cuda")
        outputs = model.generate(
            input_ids=inputs,
            max_new_tokens=512,
            temperature=0.6,
            top_p=0.95,
            do_sample=True,
            repetition_penalty=1.1,
        )
        text = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
        print(f"\nQ: {p[:80]}...\nA: {text[:400]}...")

# merge_lora.py
from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="./outputs/r1-lora-v1",
    max_seq_length=2048,
    load_in_4bit=False,  # 合并时必须用 fp16/bf16
)

# 合并并保存为 16bit
model.save_pretrained_merged(
    "./merged/r1-7b-domain-v1",
    tokenizer,
    save_method="merged_16bit",
)

# 也可以直接导出 GGUF（内部调用 llama.cpp）
model.save_pretrained_gguf(
    "./gguf/r1-7b-domain-v1",
    tokenizer,
    quantization_method="q4_k_m",
)