gleamlm 0.2.0__py3-none-any.whl

gleamlm/__init__.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from typing import Optional, Tuple
+from .models.model import GleamLMModel
+
+
+def load_model_for_inference(model_path: str, device: str = 'cuda',
+                             checkpoint: Optional[dict] = None
+                             ) -> Tuple[GleamLMModel, dict]:
+    """从 checkpoint 加载模型用于推理/评估。"""
+    import torch
+
+    if checkpoint is None:
+        try:
+            checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+        except TypeError:
+            checkpoint = torch.load(model_path, map_location=device)
+
+    if 'args' in checkpoint:
+        args = checkpoint['args']
+        config = {
+            'vocab_size': args.vocab_size,
+            'd_model': args.d_model,
+            'num_layers': args.num_layers,
+            'num_heads': args.num_heads,
+            'num_kv_heads': getattr(args, 'num_kv_heads', args.num_heads),
+            'd_ff': args.d_ff,
+            'dropout': 0.0,
+            'max_seq_len': args.max_seq_len,
+            'pad_token_id': getattr(args, 'pad_token_id', 0),
+            'tie_weights': getattr(args, 'tie_weights', True),
+            'use_flash_attn': getattr(args, 'use_flash_attn', False),
+        }
+    elif 'config' in checkpoint:
+        config = checkpoint['config']
+    else:
+        raise ValueError(
+            "Checkpoint 缺少模型结构信息。"
+            "请确保 checkpoint 包含 'args' 或 'config' 字段。"
+        )
+
+    config['dropout'] = 0.0
+
+    model = GleamLMModel(**config).to(device)
+
+    if 'model_state_dict' in checkpoint:
+        state_dict = checkpoint['model_state_dict']
+        state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+        missing, unexpected = model.load_state_dict(state_dict, strict=False)
+        if missing:
+            print(f"Warning: missing keys in checkpoint: {missing}")
+        if unexpected:
+            print(f"Warning: unexpected keys in checkpoint: {unexpected}")
+
+    model.eval()
+
+    if checkpoint.get('dtype') == 'float16':
+        model = model.half()
+
+    return model, config

gleamlm/dataset/__init__.py

@@ -0,0 +1 @@
+from .dataset import LMDataset, collate_fn

gleamlm/dataset/dataset.py

@@ -0,0 +1,127 @@
+"""GleamLM 数据集。滑动窗口切分 + 分词 + 动态 padding + numpy memmap。"""
+from __future__ import annotations
+
+import sys
+from typing import List, Tuple
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+import os
+import random
+from tqdm import tqdm
+
+from gleamlm.tokenizer.tokenizer import BBPETokenizer
+
+
+class LMDataset(Dataset):
+    """LM 数据集（memmap 版本）。分词后存磁盘，按索引切片"""
+    def __init__(self, data_dir: str, tokenizer: BBPETokenizer,
+                 max_seq_len: int, split: str = "train",
+                 stride: int | None = None,
+                 max_chars: int | None = None,
+                 ids_prefix: str = "", augment: bool = True) -> None:
+        self.tokenizer = tokenizer
+        self.max_seq_len = max_seq_len
+        self.stride = stride or max_seq_len * 3 // 4
+        self.augment = augment
+
+        text_file = os.path.join(data_dir, f"{split}.txt")
+        id_prefix = f"{ids_prefix}_" if ids_prefix else ""
+        ids_file = os.path.join(data_dir, f"{split}_{id_prefix}ids.npy")
+
+        if not os.path.exists(text_file):
+            raise FileNotFoundError(
+                f"Data file not found: {text_file}\n"
+                f"Please run data_tools/build_dataset.py first to prepare the data."
+            )
+
+        if os.path.exists(ids_file):
+            print(f"Loading pre-tokenized {split} data from {ids_file}...")
+            self.all_ids = np.load(ids_file, mmap_mode='r')
+            self.total_tokens = len(self.all_ids)
+            print(f"Loaded {split} data: {self.total_tokens} tokens")
+        else:
+            import torch.distributed as dist
+            rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+            if rank == 0:
+                with open(text_file, 'r', encoding='utf-8') as f:
+                    text = f.read() if max_chars is None else f.read(max_chars)
+                total_chars = len(text)
+
+                chunk_size = 256 * 1024
+                all_ids: list[int] = []
+                n_chunks = (total_chars + chunk_size - 1) // chunk_size
+
+                print(f"Tokenizing {split} data ({total_chars/1e9:.2f}B chars, "
+                      f"{n_chunks} chunks)...")
+                sys.stdout.flush()
+                status_path = os.path.join(data_dir, f".tokenizing_{split}")
+                for i in tqdm(range(0, total_chars, chunk_size), desc=f"  {split}",
+                              file=sys.stdout):
+                    chunk = text[i:i + chunk_size]
+                    ids = tokenizer.encode(chunk, add_bos=False, add_eos=False)
+                    all_ids.extend(ids)
+                    if (i // chunk_size) % 10 == 0:
+                        with open(status_path, "w") as sf:
+                            sf.write(f"{i+chunk_size}/{total_chars}")
+
+                print(f"  Done: {len(all_ids)} tokens")
+
+                ids_array = np.array(all_ids, dtype=np.uint32)
+                np.save(ids_file, ids_array)
+                del all_ids, text
+                print(f"  Saved to {ids_file}")
+
+                try:
+                    os.remove(status_path)
+                except OSError:
+                    pass
+
+            if dist.is_available() and dist.is_initialized():
+                dist.barrier()
+
+            self.all_ids = np.load(ids_file, mmap_mode='r')
+            self.total_tokens = len(self.all_ids)
+            print(f"Loaded {split} data: {self.total_tokens} tokens")
+
+        self.num_samples = max(0, (self.total_tokens - self.max_seq_len - 1) // self.stride + 1)
+        if self.num_samples == 0:
+            print(f"*** WARNING: 0 samples for {split}! "
+                  f"total_tokens={self.total_tokens}, max_seq_len={self.max_seq_len}. "
+                  f"Data may be too small or max_seq_len too large.")
+        print(f"Created {self.num_samples} samples for {split}")
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        start = idx * self.stride
+        end = start + self.max_seq_len + 1
+        ids = self.all_ids[start:end].astype(np.int64)
+        tensor = torch.from_numpy(ids)
+
+        if self.augment and random.random() < 0.10:
+            min_len = self.max_seq_len // 2
+            trunc = random.randint(min_len, self.max_seq_len)
+            tensor = tensor[:trunc]
+
+        return tensor
+
+
+def collate_fn(batch: List[torch.Tensor], pad_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Padding 到 batch 内最大长度，右移一位拆分为 input 和 target"""
+    max_len = max(len(sample) for sample in batch)
+
+    padded = []
+    for sample in batch:
+        if len(sample) < max_len:
+            padding = torch.full((max_len - len(sample),), pad_id, dtype=torch.long)
+            sample = torch.cat([sample, padding])
+        padded.append(sample)
+
+    batch_tensor = torch.stack(padded)
+
+    input_ids = batch_tensor[:, :-1]
+    target_ids = batch_tensor[:, 1:]
+
+    return input_ids, target_ids

gleamlm/evaluation/__init__.py

@@ -0,0 +1,11 @@
+"""GleamLM 统一评估框架 — PPL / 知识探针 / CEVAL / CMMLU"""
+
+from .ppl import evaluate_ppl, evaluate_multiple, PPLResult
+from .knowledge import evaluate_knowledge, KnowledgeResult
+from .benchmark import evaluate_ceval, evaluate_cmmlu, BenchmarkResult
+
+__all__ = [
+    "evaluate_ppl", "evaluate_multiple", "PPLResult",
+    "evaluate_knowledge", "KnowledgeResult",
+    "evaluate_ceval", "evaluate_cmmlu", "BenchmarkResult",
+]

gleamlm/evaluation/benchmark.py

@@ -0,0 +1,185 @@
+"""基准测试评估 — CEVAL / CMMLU / 多选题"""
+from __future__ import annotations
+
+import os, json, re
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn.functional as F
+
+from gleamlm.utils.torch_utils import safe_autocast
+from gleamlm.models.model import GleamLMModel
+from gleamlm.tokenizer.tokenizer import BBPETokenizer
+
+
+@dataclass
+class BenchmarkResult:
+    """基准测试结果"""
+    name: str                      # ceval / cmmlu
+    accuracy: float
+    total: int
+    correct: int
+    subject_scores: Dict[str, float] = field(default_factory=dict)
+    extra: Dict = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "benchmark": self.name,
+            "accuracy": round(self.accuracy, 4),
+            "total": self.total,
+            "correct": self.correct,
+            "subjects": self.subject_scores,
+            **self.extra,
+        }
+
+    def __repr__(self) -> str:
+        return (f"BenchmarkResult({self.name}: "
+                f"acc={self.accuracy:.2%}, {self.correct}/{self.total})")
+
+
+def _mc_generate(model: GleamLMModel, tokenizer: BBPETokenizer,
+                 prompt: str, choices: List[str],
+                 device: str) -> str:
+    """多选题生成：预填充 prompt KV Cache，仅对选项部分增量推理。
+
+    将 prompt 预填充一次得到 KV Cache，然后每个选项仅输入选项 token
+    并复用 Cache，避免重复计算 prompt 部分。
+    """
+    prompt_ids = tokenizer.encode(prompt)
+    prompt_t = torch.tensor([prompt_ids], device=device)
+
+    with torch.no_grad(), safe_autocast():
+        _, prompt_kv = model(prompt_t)
+
+    choice_scores = []
+
+    for choice in choices:
+        choice_ids = tokenizer.encode(choice)
+        if not choice_ids:
+            choice_scores.append(float('-inf'))
+            continue
+
+        choice_t = torch.tensor([choice_ids], device=device)
+
+        with torch.no_grad(), safe_autocast():
+            logits, _ = model(choice_t, past_kv_list=prompt_kv)
+
+        total_log_prob = 0.0
+        for i, token_id in enumerate(choice_ids):
+            pred_logits = logits[0, i, :]
+            log_prob = torch.nn.functional.log_softmax(pred_logits, dim=-1)[token_id].item()
+            total_log_prob += log_prob
+
+        avg_log_prob = total_log_prob / len(choice_ids)
+        choice_scores.append(avg_log_prob)
+
+    best_idx = choice_scores.index(max(choice_scores))
+    return chr(ord('A') + best_idx) if best_idx < 4 else str(best_idx)
+
+
+def _load_ceval(data_dir: str) -> List[Dict]:
+    """加载 CEVAL 数据集。支持单文件 ceval.json 或按学科拆分。"""
+    ceval_file = os.path.join(data_dir, "ceval.json")
+    if os.path.exists(ceval_file):
+        with open(ceval_file, 'r', encoding='utf-8') as f:
+            return json.load(f)
+
+    # 按学科拆分：ceval/high_school_math.json, etc.
+    items = []
+    for fname in sorted(os.listdir(data_dir)):
+        if fname.endswith('.json') and fname != 'ceval.json':
+            path = os.path.join(data_dir, fname)
+            with open(path, 'r', encoding='utf-8') as f:
+                subject_items = json.load(f)
+                subject = os.path.splitext(fname)[0]
+                for item in subject_items:
+                    item.setdefault('subject', subject)
+                items.extend(subject_items)
+    return items
+
+
+def _build_prompt(item: Dict) -> str:
+    """构建多选题 prompt: 问题 + A/B/C/D 选项"""
+    question = item.get('question', '')
+    choices = []
+    for key in ['A', 'B', 'C', 'D']:
+        if key in item:
+            choices.append(item[key])
+
+    prompt = f"{question}\n"
+    for i, choice in enumerate(choices):
+        prompt += f"{chr(ord('A') + i)}. {choice}\n"
+    prompt += "答案："
+    return prompt
+
+
+def evaluate_ceval(model: GleamLMModel, tokenizer: BBPETokenizer,
+                   data_dir: str, device: str = "cuda",
+                   subjects: Optional[List[str]] = None,
+                   max_samples: Optional[int] = None,
+                   verbose: bool = True) -> BenchmarkResult:
+    """CEVAL 中文综合能力评估。
+
+    数据集格式: [{question, A, B, C, D, answer, subject}, ...]
+    """
+    items = _load_ceval(data_dir)
+    if not items:
+        raise FileNotFoundError(f"No CEVAL data found in {data_dir}")
+
+    if subjects:
+        items = [it for it in items if it.get('subject', '') in subjects]
+    if max_samples:
+        items = items[:max_samples]
+
+    model.eval()
+    subject_correct = {}
+    subject_total = {}
+    total_correct = 0
+
+    for item in items:
+        subject = item.get('subject', 'general')
+        choices = [item.get(k, '') for k in ['A', 'B', 'C', 'D'] if k in item]
+        prompt = _build_prompt(item)
+        predicted = _mc_generate(model, tokenizer, prompt, choices, device)
+        expected = item.get('answer', '')
+
+        correct = predicted == expected
+        if correct:
+            total_correct += 1
+            subject_correct[subject] = subject_correct.get(subject, 0) + 1
+        subject_total[subject] = subject_total.get(subject, 0) + 1
+
+        if verbose and len(items) <= 20:
+            mark = "✓" if correct else "✗"
+            print(f"  [{mark}] {subject}: pred={predicted}, ans={expected} | {item.get('question', '')[:50]}...")
+
+    # 计算各学科分数
+    subject_scores = {}
+    for subj in subject_total:
+        subject_scores[subj] = round(subject_correct.get(subj, 0) / max(1, subject_total[subj]), 4)
+
+    accuracy = total_correct / max(1, len(items))
+    result = BenchmarkResult(
+        name="ceval",
+        accuracy=accuracy,
+        total=len(items),
+        correct=total_correct,
+        subject_scores=subject_scores,
+    )
+
+    if verbose:
+        print(f"\nCEVAL: {result.correct}/{result.total} = {accuracy:.2%}")
+        if len(subject_scores) > 1:
+            for subj, acc in sorted(subject_scores.items()):
+                print(f"  {subj}: {acc:.2%}")
+    return result
+
+
+def evaluate_cmmlu(model: GleamLMModel, tokenizer: BBPETokenizer,
+                   data_dir: str, device: str = "cuda",
+                   **kwargs: Any) -> BenchmarkResult:
+    """CMMLU 评估 — 同 CEVAL 格式，可直接复用。"""
+    result = evaluate_ceval(model, tokenizer, data_dir, device, **kwargs)
+    result.name = "cmmlu"
+    return result

gleamlm/evaluation/knowledge.py

@@ -0,0 +1,212 @@
+"""知识探针评估 — 填空测试 + 实体一致性检查
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import random
+from dataclasses import dataclass, field
+
+from gleamlm.utils.torch_utils import safe_autocast
+from gleamlm.models.model import GleamLMModel
+from gleamlm.tokenizer.tokenizer import BBPETokenizer
+from gleamlm.inference.sampler import sample_token
+
+
+@dataclass
+class KnowledgeResult:
+    correct: int = 0
+    wrong: int = 0
+    hallucination: int = 0
+    detailed: List[Dict] = field(default_factory=list)
+    entity_consistency: Dict[str, int] = field(default_factory=dict)
+
+    @property
+    def total(self):
+        return self.correct + self.wrong + self.hallucination
+
+    @property
+    def accuracy(self):
+        return self.correct / max(1, self.total)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "accuracy": round(self.accuracy, 4),
+            "correct": self.correct,
+            "wrong": self.wrong,
+            "hallucination": self.hallucination,
+            "total": self.total,
+            "entity_consistency": self.entity_consistency,
+            "n_consistent": sum(1 for v in self.entity_consistency.values() if v >= 2),
+        }
+
+    def __repr__(self) -> str:
+        return (f"KnowledgeResult(acc={self.accuracy:.2%}, "
+                f"{self.correct}/{self.total} correct, "
+                f"{self.hallucination} hallucinations)")
+
+
+# 默认知识探针池
+DEFAULT_FACT_PROMPTS = [
+    ("世界上最高的山峰是", "珠穆朗玛峰"),
+    ("水的化学式是", "H2O"),
+    ("中国的首都是", "北京"),
+    ("爱因斯坦提出了", "相对论"),
+    ("光的速度大约是每秒", "30万公里"),
+    ("地球绕着什么转", "太阳"),
+    ("人类登上月球是在", "1969"),
+    ("DNA的全称是", "脱氧核糖核酸"),
+    ("第一个进入太空的人是", "加加林"),
+    ("太阳系有几大行星", "八"),
+    ("诺贝尔奖的创立者是", "诺贝尔"),
+    ("光合作用需要", "阳光"),
+    ("地球的卫星是", "月球"),
+    ("世界上最大的海洋是", "太平洋"),
+    ("一年有多少天", "365"),
+    ("人体最大的器官是", "皮肤"),
+    ("鲨珑GleamLM是什么", "轻量级开源对话模型"),
+]
+
+
+DEFAULT_ENTITY_PROBES = {
+    "北京": ["中国的首都是哪里", "天安门在哪个城市", "故宫位于"],
+    "光速": ["光的速度是多少", "真空中最快的是什么"],
+    "爱因斯坦": ["爱因斯坦提出了什么", "相对论是谁提出的", "E=mc^2是谁的公式"],
+    "太阳": ["太阳是什么", "地球绕着什么转", "离我们最近的恒星是"],
+    "DNA": ["DNA是什么", "遗传物质叫什么"],
+    "诺贝尔": ["诺贝尔奖是谁设立的", "诺贝尔发明了什么"],
+    "太平洋": ["最大的海洋", "太平洋在哪"],
+    "月球": ["月球的别称", "地球的卫星是什么"],
+}
+
+
+def _simple_generate(model: GleamLMModel, tokenizer: BBPETokenizer,
+                     prompt: str, device: str,
+                     max_new_tokens: int = 64, temperature: float = 0.7,
+                     top_k: int = 50) -> str:
+    """独立生成函数 — 使用 KV Cache 逐 token 增量推理。
+
+    首步预填充 prompt，后续每步仅输入新 token 并复用 KV Cache，
+    RoPE 通过 offset 自动递增位置编码。
+    """
+    prompt_ids = tokenizer.encode(prompt)
+    input_ids = torch.tensor([prompt_ids], device=device)
+    generated_ids = prompt_ids.copy()
+    past_kv = None
+
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            with safe_autocast():
+                logits, past_kv = model(input_ids, past_kv_list=past_kv)
+            next_token = sample_token(
+                logits[:, -1, :],
+                temperature=temperature,
+                top_k=top_k,
+                top_p=0.0,
+                generated_ids=generated_ids,
+            )
+            token_id = next_token.item()
+            if token_id in {tokenizer.eos_id, tokenizer.special_tokens.get("<|im_end|>")}:
+                break
+            generated_ids.append(token_id)
+            input_ids = torch.tensor([[token_id]], device=device)
+
+    full = tokenizer.decode(generated_ids)
+    if full.startswith(prompt):
+        generated = full[len(prompt):].strip()
+    else:
+        generated = full.strip()
+    return generated[:200]
+
+
+def _check_answer(generated: str, expected: str,
+                  hallucination_keywords: Optional[List[str]] = None) -> str:
+    """检查生成文本是否包含预期答案。返回 CORRECT / WRONG / HALLUCINATION"""
+    gen_lower = generated.lower().replace(" ", "")
+    answers = [a.strip().lower() for a in expected.split(",")]
+    for ans in answers:
+        if ans in gen_lower:
+            return "CORRECT"
+    if hallucination_keywords:
+        for hw in hallucination_keywords:
+            if hw in generated:
+                return "HALLUCINATION"
+    return "WRONG"
+
+
+def evaluate_knowledge(model: GleamLMModel, tokenizer: BBPETokenizer,
+                       device: str = "cuda",
+                       fact_prompts: Optional[List[Tuple[str, str]]] = None,
+                       entity_probes: Optional[Dict[str, List[str]]] = None,
+                       hallucination_keywords: Optional[List[str]] = None,
+                       max_new_tokens: int = 64,
+                       temperature: float = 0.7,
+                       verbose: bool = True) -> KnowledgeResult:
+    """运行知识探针评估。"""
+    if fact_prompts is None:
+        fact_prompts = DEFAULT_FACT_PROMPTS
+    if entity_probes is None:
+        entity_probes = DEFAULT_ENTITY_PROBES
+
+    model.eval()
+    result = KnowledgeResult()
+
+    # A1: 事实填空
+    if verbose:
+        print(f"\n{'='*60}")
+        print(f"A1: FACT FILL-IN ({len(fact_prompts)} prompts)")
+        print(f"{'='*60}")
+
+    for prompt, expected in fact_prompts:
+        generated = _simple_generate(model, tokenizer, prompt, device,
+                                     max_new_tokens, temperature)
+        label = _check_answer(generated, expected, hallucination_keywords)
+
+        if label == "CORRECT":
+            result.correct += 1
+        elif label == "HALLUCINATION":
+            result.hallucination += 1
+        else:
+            result.wrong += 1
+
+        result.detailed.append({
+            "prompt": prompt, "expected": expected,
+            "generated": generated[:100], "result": label
+        })
+
+        if verbose:
+            symbols = {"CORRECT": "OK", "HALLUCINATION": "HALU", "WRONG": "WRONG"}
+            print(f"  [{symbols[label]:>4}] {prompt} → {generated[:60]}...")
+
+    if verbose:
+        print(f"\n  Accuracy: {result.correct}/{result.total} ({result.accuracy:.1%}), "
+              f"{result.hallucination} hallucinations")
+
+    # A2: 实体探针
+    if verbose:
+        print(f"\n{'='*60}")
+        print(f"A2: ENTITY PROBE ({len(entity_probes)} entities)")
+        print(f"{'='*60}")
+
+    for entity, questions in entity_probes.items():
+        correct = 0
+        for q in questions:
+            generated = _simple_generate(model, tokenizer, q, device,
+                                        max_new_tokens, temperature)
+            if entity.lower() in generated.lower():
+                correct += 1
+        result.entity_consistency[entity] = correct
+
+        if verbose:
+            icon = "***" if correct == len(questions) else ("*" if correct > 0 else " ")
+            print(f"  [{icon}] {entity}: {correct}/{len(questions)}")
+
+    if verbose:
+        consistent = sum(
+            1 for entity, v in result.entity_consistency.items()
+            if entity in entity_probes and v >= len(entity_probes[entity]) * 0.67
+        )
+        print(f"\n  Consistent entities: {consistent}/{len(entity_probes)}")
+
+    return result