MedicalGPT 学习指北

@torch.no_grad()
def calc_ppl_sliding_window(
    model,
    tokenizer,
    texts: list[str],
    max_length: int = 2048,
    stride: int = 512,
    batch_size: int = 1,
    device: str = "cuda",
    verbose: bool = True,
) -> dict:
    model.eval()
    loss_fn = CrossEntropyLoss(reduction="sum")

    total_nll = 0.0
    total_tokens = 0
    per_text_ppls = []

    iterator = tqdm(texts, desc="Computing PPL") if verbose else texts

    for text in iterator:
        encodings = tokenizer(text, return_tensors="pt", truncation=False)
        input_ids = encodings.input_ids.to(device)  # [1, seq_len]
        seq_len = input_ids.size(1)

        text_nll = 0.0
        text_tokens = 0
        prev_end = 0

        for begin in range(0, seq_len, stride):
            end = min(begin + max_length, seq_len)
            # 当前窗口的 input_ids
            chunk = input_ids[:, begin:end]
            # 只计算新 token 的 loss（去掉与上一个窗口重叠的部分）
            target_len = end - prev_end

            with torch.no_grad():
                outputs = model(chunk)
                logits = outputs.logits  # [1, chunk_len, vocab_size]

            # 标准 next-token shift：logits[i] 预测 token[i+1]
            # shift 后长度为 chunk_len - 1
            shift_logits = logits[0, :-1, :]  # [chunk_len-1, vocab]
            shift_labels = chunk[0, 1:]  # [chunk_len-1]

            # 只对新 token 部分（最后 target_len 个位置）计算 loss
            # 同时 clamp 防止 target_len > shift 后的长度（极短文本边界情况）
            actual_target = min(target_len, shift_labels.size(0))
            shift_logits = shift_logits[-actual_target:].contiguous()
            shift_labels = shift_labels[-actual_target:].contiguous()

            nll = loss_fn(shift_logits, shift_labels)
            text_nll += nll.item()
            text_tokens += target_len

            prev_end = end
            if end == seq_len:
                break

        if text_tokens > 0:
            text_ppl = math.exp(text_nll / text_tokens)
            per_text_ppls.append(text_ppl)
            total_nll += text_nll
            total_tokens += text_tokens

    overall_ppl = (
        math.exp(total_nll / total_tokens) if total_tokens > 0 else float("inf")
    )

    return {
        "ppl": round(overall_ppl, 4),
        "total_tokens": total_tokens,
        "num_texts": len(per_text_ppls),
        "per_text_ppls": [round(p, 4) for p in per_text_ppls],
        "avg_text_ppl": round(sum(per_text_ppls) / len(per_text_ppls), 4)
        if per_text_ppls
        else None,
    }

python pretraining.py \
    --model_name_or_path /root/autodl-tmp/models/Qwen/Qwen3___5-2B \
    --train_file_dir /root/MedicalGPT/data/pretrain/train \
    --validation_file_dir /root/MedicalGPT/data/pretrain/validate \
    --per_device_train_batch_size 16 \
    --per_device_eval_batch_size 4 \
    --do_train \
    --do_eval \
    --use_peft True \
    --seed 42 \
    --num_train_epochs 0.5 \
    --learning_rate 5e-5 \
    --warmup_steps 5 \
    --weight_decay 0.01 \
    --logging_strategy steps \
    --logging_steps 10 \
    --eval_steps 50 \
    --eval_strategy steps \
    --save_steps 500 \
    --save_strategy steps \
    --save_total_limit 13 \
    --gradient_accumulation_steps 8 \
    --preprocessing_num_workers 10 \
    --block_size 512 \
    --packing True \
    --output_dir /root/autodl-tmp/models/medicalGPT/pt_lora \
    --logging_first_step True \
    --target_modules all \
    --lora_rank 8 \
    --lora_alpha 16 \
    --lora_dropout 0.05 \
    --torch_dtype bfloat16 \
    --bf16 \
    --report_to tensorboard \
    --gradient_checkpointing True \
    --cache_dir ./cache

from dataclasses import dataclass, field
from typing import Optional

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser


@dataclass
class ModelArguments:
    model_dir: str = field()
    lora_dir: str = field()
    output_dir: str = field()
    tok_dir: Optional[str] = field(default=None)

    def __post_init__(self):
        if self.tok_dir is None:
            self.tok_dir = self.model_dir


if __name__ == "__main__":
    parser = HfArgumentParser((ModelArguments))
    (args,) = parser.parse_args_into_dataclasses()

    model = AutoModelForCausalLM.from_pretrained(
        args.model_dir, torch_dtype=torch.float16, device_map="auto"
    )
    model = PeftModel.from_pretrained(model, args.lora_dir)
    model = model.merge_and_unload()
    model.save_pretrained(args.output_dir)

    tokenizer = AutoTokenizer.from_pretrained(args.tok_dir)
    tokenizer.save_pretrained(args.output_dir)

    print(f"model saved in {args.output_dir}")

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import torch
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from sklearn.metrics import accuracy_score, f1_score
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    EvalPrediction,
    HfArgumentParser,
    Trainer,
)
from transformers.utils.versions import require_version

from datasets import DatasetDict, load_dataset


@dataclass
class ModelArguments:
    model_name_or_path: str = field()
    tokenizer_name_or_path: Optional[str] = field(default=None)
    dtype: Optional[str] = field(
        default="float32",
        metadata={
            "choices": ["auto", "bfloat16", "float16", "float32"],
        },
    )
    load_in_8bit: bool = field(default=False)
    load_in_4bit: bool = field(default=False)
    trust_remote_code: Optional[bool] = field(default=False)
    use_fast_tokenizer: bool = field(default=False)
    cache_dir: Optional[str] = field(default=None)

    def __post_init__(self):

        assert not (self.load_in_4bit and self.load_in_8bit)

        if self.tokenizer_name_or_path is None:
            self.tokenizer_name_or_path = self.model_name_or_path


@dataclass
class DataArguments:
    dataset_name: Optional[str] = field(default=None)
    dataset_config_name: Optional[str] = field(default=None)
    train_dir: Optional[str] = field(default=None)
    validate_dir: Optional[str] = field(default=None)
    max_train_samples: Optional[int] = field(default=None)
    max_validate_samples: Optional[int] = field(default=None)
    block_size: Optional[int] = field(default=512)
    validation_split_percentage: Optional[float] = field(default=0.1)
    packing: bool = field(default=True)
    streaming: bool = field(default=False)

    def __post_init__(self):
        if self.streaming:
            require_version(
                "datasets>=2.0.0", "The streaming feature requires `datasets>=2.0.0`"
            )


@dataclass
class PeftArguments:
    use_peft: bool = field(default=True)
    target_modules: Optional[str] = field(default="all")
    lora_rank: Optional[int] = field(default=8)
    lora_dropout: Optional[float] = field(default=0.0)
    lora_alpha: Optional[float] = field(default=16.0)
    qlora: bool = field(default=False)


@dataclass
class TrainingArguments:
    do_train: bool = field(default=False)
    do_eval: bool = field(default=False)
    output_dir: str = field()


def compute_metrics(eval_pred: EvalPrediction) -> dict[str, float]:
    logits, labels = eval_pred
    preds = torch.argmax(logits, dim=-1)

    return {
        "accuracy": accuracy_score(labels, preds),
        "f1": f1_score(labels, preds, average="macro"),
    }


def main():
    logger = logging.getLogger(__name__)
    parser = HfArgumentParser(
        (ModelArguments, DataArguments, PeftArguments, TrainingArguments)
    )
    model_args, data_args, peft_args, train_args = parser.parse_args_into_dataclasses(
        return_remaining_strings=True
    )[:4]

    model_dtype = (
        model_args.dtype
        if model_args in ["auto", None]
        else getattr(torch, model_args.dtype)
    )

    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name_or_path,
        use_fast=model_args.use_fast_tokenizer,
        cache_dir=model_args.cache_dir,
        trust_remote_code=model_args.trust_remote_code,
    )
    tokenizer.add_eos_token = True

    if data_args.block_size is None:
        block_size = tokenizer.model_max_length
    else:
        if data_args.block_size > tokenizer.model_max_length:
            logger.warning(
                "block_size should be less then or equal to tokenizer.model_max_length"
            )
        block_size = min(data_args.block_size, tokenizer.model_max_length)

    def tokenize_padding_fn(examples: dict[str, str]):
        inputs = tokenizer(
            examples["text"],
            truncation=True,
            padding="max_length",
            max_length=block_size,
        )
        inputs["labels"] = [
            [
                input_id if input_id != tokenizer.pad_token_id else -100
                for input_id in sample
            ]
            for sample in inputs["input_ids"]
        ]
        return inputs

    def tokenize_packing_fn(examples: dict[str, str]):
        inputs = tokenizer(examples["text"], add_special_token=True)
        all_input_ids = []
        all_attn_masks = []
        for ids, masks in zip(inputs["input_ids"], inputs["attention_mask"]):
            all_input_ids.extend(ids)
            all_attn_masks.extend(masks)

        num_blocks = len(all_input_ids) // block_size

        output = {
            "input_ids": [
                all_input_ids[i * block_size : (i + 1) * block_size]
                for i in range(num_blocks)
            ],
            "attention_mask": [
                all_attn_masks[i * block_size : (i + 1) * block_size]
                for i in range(num_blocks)
            ],
        }
        output["labels"] = output["input_ids"].copy()
        return output

    model_kwargs = {}
    if model_args.load_in_4bit:
        if peft_args.use_qlora:
            model_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=model_dtype,
                bnb_4bit_use_double_quant=True,
            )
        else:
            model_kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=model_dtype,
            )
    elif model_args.load_in_8bit:
        model_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)

    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        dtype=model_dtype,
        cache_dir=model_args.cache_dir,
        trust_remote_code=model_args.trust_remote_code,
        **model_kwargs,
    )

    if model_args.load_in_4bit or model_args.load_in_8bit:
        model = prepare_model_for_kbit_training(model)

    if peft_args.use_peft:
        model = get_peft_model(
            model,
            LoraConfig(
                r=peft_args.lora_rank,
                lora_alpha=peft_args.lora_alpha,
                lora_dropout=peft_args.lora_dropout,
                bias="none",
                task_type=TaskType.CAUSAL_LM,
                target_modules="all-linear",
            ),
        )
        logger.info("use lora training")
    else:
        logger.info("use full-parameter training")

    dataset = DatasetDict()
    if data_args.dataset_name:
        logger.info("begin downloading datasets")
        raw_dataset = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            streaming=data_args.streaming,
        )
    else:
        data_files = {}
        if data_args.train_dir and (folder := Path(data_args.train_dir)).exists():
            data_files["train"] = [
                str(p) for p in folder.rglob("*") if p.suffix in {".json", ".jsonl"}
            ]
        if data_args.validate_dir and (folder := Path(data_args.validate_dir)).exists():
            data_files["validation"] = [
                str(p) for p in folder.rglob("*") if p.suffix in {".json", ".jsonl"}
            ]

        raw_dataset = load_dataset(
            "json",
            data_files=data_files,
            cache_dir=model_args.cache_dir,
            streaming=data_args.streaming,
        )

    if train_args.do_eval and "validation" not in raw_dataset.keys():
        dataset["train"] = raw_dataset["train"].train_test_split(
            test_size=data_args.validation_split_percentage
        )
        dataset["validation"] = dataset.pop("test")
    else:
        dataset = raw_dataset

    logger.info(f"dataset: {dataset}")

    lm_datasets = dataset.map(
        tokenize_packing_fn if data_args.packing else tokenize_padding_fn,
        batched=True,
        num_proc=data_args.preprocessing_num_workers if data_args.streaming else 1,
        remove_columns=list(dataset["train"].features()),
        desc="Running tokenizer on dataset",
    )

    train_dataset = lm_datasets["train"]
    if data_args.max_train_samples and data_args.max_train_samples < len(train_dataset):
        train_dataset = train_dataset.select(range(data_args.max_train_samples))

    val_dataset = lm_datasets["validation"]
    if data_args.max_validate_sampels and data_args.max_validate_sampels < len(
        val_dataset
    ):
        val_dataset = val_dataset.select(range(data_args.max_validate_sampels))

    trainer = Trainer(
        model=model,
        args=train_args,
        train_dataset=train_dataset if train_args.do_train else None,
        eval_dataset=val_dataset if train_args.do_eval else None,
        processing_class=tokenizer,
        compute_metrics=compute_metrics if train_args.do_eval else None,
    )

    trainer.train()
    model.save_pretrained(train_args.output_dir)


if __name__ == "__main__":
    main()

@torch.no_grad()
def calc_ppl_batch(
    model,
    tokenizer,
    texts: list[str],
    max_length: int = 2048,
    stride: int = 512,
    batch_size: int = 2,
    device: str = "cuda",
):
    model.eval()
    loss_fn = CrossEntropyLoss(reduction="none")

    all_ppls = []

    for i in tqdm(range(0, len(texts), batch_size), desc="PPL"):
        batch_texts = texts[i : i + batch_size]

        enc = tokenizer(
            batch_texts,
            return_tensors="pt",
            padding=True,
            truncation=False,
        )

        input_ids = enc.input_ids.to(device)
        attention_mask = enc.attention_mask.to(device)

        batch_size_, seq_len = input_ids.shape

        # 每条样本独立统计
        nll_sum = torch.zeros(batch_size_, device=device)
        token_count = torch.zeros(batch_size_, device=device)

        prev_end = 0

        for begin in range(0, seq_len, stride):
            end = min(begin + max_length, seq_len)

            chunk_ids = input_ids[:, begin:end]
            chunk_mask = attention_mask[:, begin:end]

            outputs = model(chunk_ids)
            logits = outputs.logits

            shift_logits = logits[:, :-1, :]
            shift_labels = chunk_ids[:, 1:]
            shift_mask = chunk_mask[:, 1:]

            # 计算逐 token loss
            loss = loss_fn(
                shift_logits.reshape(-1, shift_logits.size(-1)),
                shift_labels.reshape(-1),
            )
            loss = loss.view(batch_size_, -1)

            # mask padding
            loss = loss * shift_mask

            # sliding window 只统计新 token
            target_len = end - prev_end
            loss = loss[:, -target_len:]
            mask = shift_mask[:, -target_len:]

            nll_sum += loss.sum(dim=1)
            token_count += mask.sum(dim=1)

            prev_end = end
            if end == seq_len:
                break

        ppl = torch.exp(nll_sum / token_count)
        all_ppls.extend(ppl.detach().cpu().tolist())

    return all_ppls

def vectorize_source(
    self,
    samples: list[dict],
    shard_size: int = 50000,
    shard_dir: str = ".cache/shard_cache",
    ans_truncate: int = 200,
    shard_info_path: str | None = None,
):
    shard_cache = Path(shard_dir)
    shard_cache.mkdir(parents=True, exist_ok=True)
    n_samples = len(samples)
    shard_info = []
    for shard_id, i in enumerate(range(0, n_samples, shard_size)):
        emb_dir = shard_cache / f"shard_{shard_id}.npy"
        meta_dir = shard_cache / f"shard_{shard_id}.jsonl"
        batch = samples[i : i + shard_size]
        if emb_dir.exists() and meta_dir.exists():
            tqdm.write(f"Shard {shard_id} 已缓存")
        else:
            texts = [
                item.get("instruction", "") + item.get("output", "")[:ans_truncate]
                for item in batch
            ]
            embs = self.embed(texts, desc=f"Shard {shard_id} Embedding")
            np.save(str(emb_dir), embs)
            with open(meta_dir, "w", encoding="utf-8") as f:
                for sample in batch:
                    f.write(json.dumps(sample, ensure_ascii=False) + "\n")
        shard_info.append(
            {
                "emb_dir": str(emb_dir),
                "meta_dir": str(meta_dir),
                "n_shard": len(batch),
                "offset": i,
            }
        )
    final_shard_info_path = (
        Path(shard_info_path)
        if shard_info_path is not None
        else shard_cache / "shard_info.json"
    )
    final_shard_info_path.parent.mkdir(parents=True, exist_ok=True)
    with open(final_shard_info_path, "w", encoding="utf-8") as f:
        json.dump(shard_info, f, ensure_ascii=False, indent=2)
    tqdm.write(f"Shard info saved: {final_shard_info_path}")
    return shard_info

def _search_one_shard(self, shard: dict):
    emb_path = Path(shard["emb_dir"])
    offset = int(shard["offset"])
    src_vec = np.load(emb_path)
    if src_vec.ndim != 2:
        raise ValueError(
            f"source shard must be 2D, got {emb_path}: {src_vec.shape}"
        )
    src_vec = src_vec.astype("float32")
    if src_vec.shape[0] == 0:
        return None
    # 维度一致性检查
    if src_vec.shape[1] != self.anc_vec.shape[1]:
        raise ValueError(
            f"dim mismatch: anchor_dim={self.anc_vec.shape[1]}, "
            f"source_dim={src_vec.shape[1]}, file={emb_path}"
        )
    if self.normalize:
        src_vec = self.normalize_vectors(src_vec)
    dim = src_vec.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(src_vec)
    k_eff = min(self.top_k, src_vec.shape[0])
    if k_eff <= 0:
        return None
    scores, local_ids = index.search(self.anc_vec, k_eff)
    return scores, local_ids, offset

def query_topk(self) -> list[list[tuple[float, int]]]:
    n_anchor = self.anc_vec.shape[0]
    # 小顶堆，存 (score, source_global_id)
    heaps: list[list[tuple[float, int]]] = [[] for _ in range(n_anchor)]
    for shard in tqdm(self.shard_info, desc="Searching shards"):
        out = self._search_one_shard(shard)
        if out is None:
            continue
        scores, local_ids, offset = out
        # 遍历每个 anchor 的局部 topK
        for a_id in range(n_anchor):
            for j in range(local_ids.shape[1]):
                lid = int(local_ids[a_id, j])
                if lid < 0:
                    continue
                gid = offset + lid
                sc = float(scores[a_id, j])
                heap = heaps[a_id]
                if len(heap) < self.top_k:
                    heapq.heappush(heap, (sc, gid))
                else:
                    if sc > heap[0][0]:
                        heapq.heapreplace(heap, (sc, gid))

python supervised_finetuning.py \
    --model_name_or_path /root/autodl-tmp/models/Qwen/Qwen3___5-2B \
    --train_file_dir ./data/finetune \
    --per_device_train_batch_size 12 \
    --do_train \
    --do_eval \
    --template_name qwen \
    --use_peft True \
    --model_max_length 512 \
    --max_eval_samples 10 \
    --num_train_epochs 1 \
    --learning_rate 2e-5 \
    --warmup_steps 5 \
    --weight_decay 0.05 \
    --logging_strategy steps \
    --logging_steps 10 \
    --eval_steps 50 \
    --eval_strategy steps \
    --save_steps 500 \
    --save_strategy steps \
    --save_total_limit 13 \
    --gradient_accumulation_steps 2 \
    --preprocessing_num_workers 4 \
    --output_dir /root/autodl-tmp/models/medicalGPT/sft_on_base_better_dataset \
    --logging_first_step True \
    --target_modules all \
    --lora_rank 8 \
    --lora_alpha 16 \
    --lora_dropout 0.0 \
    --torch_dtype bfloat16 \
    --bf16 \
    --report_to tensorboard \
    --gradient_checkpointing True \
    --cache_dir ./cache --flash_attn True

import json

SYSTEM_PROMPT = """You are a medical QA evaluator. Score each answer independently based on the reference answer.

## Scoring Criteria (0.0 - 1.0 continuous scale)
- 1.0: Fully correct and complete. Key medical facts match the reference precisely.
- 0.8-0.9: Mostly correct with very minor omissions or slight inaccuracies.
- 0.6-0.7: Generally correct but missing important details or contains some inaccuracies.
- 0.4-0.5: Partially correct. Core idea is present but significant details are missing or some errors exist.
- 0.2-0.3: Mostly incorrect. Only small parts are relevant or correct.
- 0.0-0.1: Completely incorrect, irrelevant, or refusal to answer.

## Important Rules
- Judge based on medical accuracy, not writing style or length
- If the answer would be harmful to a patient, always score 0.0
- Use the full range of scores (not only round numbers like 0.0, 0.5, 1.0)

## Output Rules (STRICT)
- Output MUST be a valid JSON array of floats
- DO NOT include any explanation, text, or markdown
- DO NOT include trailing commas
- The length of the array MUST equal the number of input items
- Each score MUST be a float between 0.0 and 1.0
- Keep at most 2 decimal places

## Enforcement
- You must strictly follow the scoring rubric
- Your output will be automatically parsed; any format error will be considered a failure
- Do not output anything other than the JSON array
"""


def build_input(prompts, references, answers) -> str:
    tmp_lines = []
    for idx, (prompt, refer, ans) in enumerate(zip(prompts, references, answers)):
        tmp_lines.append(
            f"QA {idx + 1}\n"
            + f"Question: {prompt}\n"
            + f"Reference Answer: {refer}\n"
            + f"Answer: {ans}\n"
        )
    input_text = "".join(tmp_lines)
    return f"""Score the following QA pairs based on the reference answers.
## Input
{input_text}
"""


def wrapper(judge_model, tokenizer):
    def accuracy_reward(prompts, completions, **kwargs):
        references = kwargs.get("reference")
        answers = [completion[0]["content"] for completion in completions]
        assert len(prompts) == len(references) == len(answers)

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": build_input(prompts, references, answers)},
        ]
        print(messages)
        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
            add_generation_prompt=True,
            enable_thinking=False,
            use_thinking=False,
        ).to(device)

        generated_ids = judge_model.generate(
            inputs=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            do_sample=False,
            max_new_tokens=8 * len(prompts) + 16,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
        response = tokenizer.decode(
            generated_ids[0][len(inputs["input_ids"][0]) :], skip_special_tokens=True
        )
        try:
            scores = json.loads(response)
        except json.JSONDecodeError:
            scores = [0.0] * len(prompts)

        assert isinstance(scores, list)
        return scores

    return accuracy_reward