初始化项目，由ModelHub XC社区提供模型

Model: pathcosmos/frankenstallm Source: Original Platform
2026-07-14 04:21:16 +08:00
commit d4abdb70fa
516 changed files with 5584477 additions and 0 deletions
--- a/source/eval/tasks/generation_task.py
+++ b/source/eval/tasks/generation_task.py
@@ -0,0 +1,472 @@
+"""
+generation_task.py — Text generation quality evaluation tasks.
+
+Top-level functions for ProcessPoolExecutor (spawn) compatibility:
+  - eval_generation(device) -> dict
+  - eval_repetition_grid(device) -> dict
+
+Helper functions (also top-level, used internally):
+  - top_p_filtering(logits, top_p, top_k)
+  - generate_one(model, tokenizer, prompt, temperature, ...)
+  - compute_ngram_rep(text, n)
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+_DEFAULT_CHECKPOINT = str(_PROJECT_ROOT / "checkpoints" / "korean_3b_fp8_run1" / "checkpoint-0057000")
+CHECKPOINT = os.environ.get("EVAL_CHECKPOINT", _DEFAULT_CHECKPOINT)
+TOKENIZER_PATH = os.environ.get("EVAL_TOKENIZER", str(_PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"))
+
+# Chat template support for SFT models
+USE_CHAT_TEMPLATE = os.environ.get("USE_CHAT_TEMPLATE", "0") == "1"
+CHAT_TEMPLATE_FMT = "<|user|>\n{prompt}\n<|assistant|>\n"
+DATA_DIR = _PROJECT_ROOT / "data"
+SEQ_LEN = 2048
+STRIDE = 512
+BATCH_SIZE = 32
+
+# ---------------------------------------------------------------------------
+# Prompt / temperature constants
+# ---------------------------------------------------------------------------
+
+PROMPTS = [
+    "대한민국의 수도는",
+    "인공지능이란",
+    "한국의 전통 음식 중에서",
+    "지구 온난화의 주요 원인은",
+    "프로그래밍을 배우려면",
+    "조선시대에는",
+    "물리학에서 에너지란",
+    "한국어는 세계에서",
+    "경제 성장을 위해서는",
+    "우주 탐사의 역사를 보면",
+    "머신러닝과 딥러닝의 차이는",
+    "한국 문학의 대표적인 작품으로는",
+    "양자 컴퓨터란",
+    "건강한 식습관을 위해서는",
+    "세계 2차 대전 이후",
+]
+
+TEMPERATURES = [0.0, 0.5, 0.8, 1.0]
+
+REP_GRID = [
+    {"name": "greedy",       "temperature": 0.0, "repetition_penalty": 1.0},
+    {"name": "t0.5",         "temperature": 0.5, "repetition_penalty": 1.0},
+    {"name": "t0.5_rep1.1",  "temperature": 0.5, "repetition_penalty": 1.1},
+    {"name": "t0.7",         "temperature": 0.7, "repetition_penalty": 1.0},
+    {"name": "t0.7_rep1.1",  "temperature": 0.7, "repetition_penalty": 1.1},
+    {"name": "t0.7_rep1.2",  "temperature": 0.7, "repetition_penalty": 1.2},
+    {"name": "t0.7_rep1.3",  "temperature": 0.7, "repetition_penalty": 1.3},
+    {"name": "t0.9",         "temperature": 0.9, "repetition_penalty": 1.0},
+    {"name": "t0.9_rep1.1",  "temperature": 0.9, "repetition_penalty": 1.1},
+    {"name": "t0.9_rep1.2",  "temperature": 0.9, "repetition_penalty": 1.2},
+    {"name": "t1.0",         "temperature": 1.0, "repetition_penalty": 1.0},
+    {"name": "t1.0_rep1.1",  "temperature": 1.0, "repetition_penalty": 1.1},
+]
+
+
+# ---------------------------------------------------------------------------
+# Shared model utilities
+# ---------------------------------------------------------------------------
+
+def _load_model(device: str):
+    """Load FRANKENSTALLM 3B from checkpoint onto the given device."""
+    from model.transformer import LLM  # type: ignore[import]
+
+    model = LLM.from_pretrained(CHECKPOINT)
+    model = model.to(device=device, dtype=torch.bfloat16)
+    model.eval()
+    return model
+
+
+def _load_tokenizer():
+    """Load the Korean SentencePiece tokenizer."""
+    from tokenizers import Tokenizer  # type: ignore[import]
+
+    return Tokenizer.from_file(TOKENIZER_PATH)
+
+
+# ---------------------------------------------------------------------------
+# Generation helpers (top-level for pickle compatibility)
+# ---------------------------------------------------------------------------
+
+def top_p_filtering(logits: torch.Tensor, top_p: float = 0.9, top_k: int = 0) -> torch.Tensor:
+    """Apply top-p (nucleus) and/or top-k filtering to a logits tensor.
+
+    Args:
+        logits: Shape (..., vocab_size).
+        top_p:  Nucleus probability threshold in (0, 1). 0 or 1 disables.
+        top_k:  Keep only the top-k tokens. 0 disables.
+
+    Returns:
+        Filtered logits tensor of the same shape.
+    """
+    if logits.dim() == 1:
+        logits = logits.unsqueeze(0)
+        squeeze = True
+    else:
+        squeeze = False
+
+    if top_k > 0:
+        k = min(top_k, logits.size(-1))
+        kth = torch.topk(logits, k, dim=-1).values[:, -1, None]
+        logits = logits.masked_fill(logits < kth, float("-inf"))
+
+    if 0.0 < top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, dim=-1, descending=True)
+        cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        remove = cum_probs - F.softmax(sorted_logits, dim=-1) >= top_p
+        sorted_logits[remove] = float("-inf")
+        logits = torch.zeros_like(logits).scatter_(-1, sorted_idx, sorted_logits)
+
+    if squeeze:
+        logits = logits.squeeze(0)
+    return logits
+
+
+def generate_one(
+    model,
+    tokenizer,
+    prompt: str,
+    temperature: float,
+    top_p: float = 0.9,
+    top_k: int = 50,
+    max_new_tokens: int = 256,
+    device: str = "cuda:0",
+    repetition_penalty: float = 1.0,
+) -> tuple[str, int, bool]:
+    """Generate a single continuation for a prompt using the given model.
+
+    Args:
+        model:              Pre-loaded language model (eval mode).
+        tokenizer:          Tokenizer with encode/decode methods.
+        prompt:             Input prompt string.
+        temperature:        Sampling temperature. 0.0 = greedy.
+        top_p:              Nucleus filtering threshold.
+        top_k:              Top-k filtering count.
+        max_new_tokens:     Maximum number of tokens to generate.
+        device:             CUDA device string.
+        repetition_penalty: Penalty > 1.0 discourages token repetition.
+
+    Returns:
+        Tuple of (generated_text, num_new_tokens, hit_eos).
+    """
+    input_ids = torch.tensor(
+        [tokenizer.encode(prompt).ids], dtype=torch.long, device=device
+    )
+    eos_id = tokenizer.token_to_id("</s>")
+    generated = input_ids
+    new_ids: list[int] = []
+    hit_eos = False
+
+    for _ in range(max_new_tokens):
+        logits_all, _ = model(generated)
+        logits = logits_all[:, -1, :].clone()
+
+        if repetition_penalty != 1.0:
+            for tid in set(generated[0].tolist()):
+                if logits[0, tid] > 0:
+                    logits[0, tid] /= repetition_penalty
+                else:
+                    logits[0, tid] *= repetition_penalty
+
+        if temperature == 0.0:
+            next_id = logits.argmax(dim=-1, keepdim=True)
+        else:
+            logits = logits / max(temperature, 1e-8)
+            logits = top_p_filtering(logits, top_p=top_p, top_k=top_k)
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+
+        generated = torch.cat([generated, next_id], dim=-1)
+        new_ids.append(next_id.item())
+
+        if eos_id is not None and next_id.item() == eos_id:
+            hit_eos = True
+            break
+
+    text = tokenizer.decode(new_ids)
+    return text, len(new_ids), hit_eos
+
+
+def compute_ngram_rep(text: str, n: int) -> float:
+    """Compute n-gram repetition rate for a whitespace-tokenized string.
+
+    Repetition rate = 1 - (unique n-grams / total n-grams).
+    A value of 0 means no repeated n-grams; 1 means all n-grams are repeated.
+
+    Args:
+        text: Input text (whitespace-tokenized).
+        n:    N-gram order (1, 2, 3, 4, ...).
+
+    Returns:
+        Float in [0, 1].
+    """
+    tokens = text.split()
+    if len(tokens) < n:
+        return 0.0
+    ngrams = [tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)]
+    if not ngrams:
+        return 0.0
+    return 1.0 - len(set(ngrams)) / len(ngrams)
+
+
+def compute_diversity_metrics(text: str) -> dict:
+    """N-gram 반복률을 보완하는 어휘 다양성 메트릭.
+
+    - Distinct-n (Li et al., 2016): 고유 n-gram 비율
+    - Type-Token Ratio: 어휘 풍부도
+    """
+    tokens = text.split()
+    n = len(tokens)
+    if n == 0:
+        return {"distinct_1": 0.0, "distinct_2": 0.0, "distinct_3": 0.0,
+                "type_token_ratio": 0.0, "vocab_size": 0, "total_tokens": 0}
+
+    unigrams = set(tokens)
+    bigrams = set(zip(tokens, tokens[1:])) if n > 1 else set()
+    trigrams = set(zip(tokens, tokens[1:], tokens[2:])) if n > 2 else set()
+
+    return {
+        "distinct_1": len(unigrams) / n,
+        "distinct_2": len(bigrams) / max(n - 1, 1),
+        "distinct_3": len(trigrams) / max(n - 2, 1),
+        "type_token_ratio": len(unigrams) / n,
+        "vocab_size": len(unigrams),
+        "total_tokens": n,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main task functions (must be top-level for pickle / spawn compatibility)
+# ---------------------------------------------------------------------------
+
+def eval_generation(device: str) -> dict:
+    """Evaluate generation quality: 15 prompts x 4 temperatures.
+
+    For each (prompt, temperature) combination:
+      - Generates up to 256 new tokens
+      - Computes 1-gram through 4-gram repetition rates
+
+    Args:
+        device: CUDA device string, e.g. "cuda:4".
+
+    Returns:
+        Dict with keys:
+          - summary: aggregate statistics across all generations
+          - samples: list of per-generation result dicts
+    """
+    torch.cuda.set_device(int(device.split(":")[-1]))
+    print(f"[GEN {device}] Loading model...")
+    model = _load_model(device)
+    tokenizer = _load_tokenizer()
+    t0 = time.time()
+
+    results: list[dict] = []
+    total_combinations = len(PROMPTS) * len(TEMPERATURES)
+    done = 0
+
+    if USE_CHAT_TEMPLATE:
+        print(f"[GEN {device}] Chat template ENABLED", flush=True)
+
+    for prompt in PROMPTS:
+        effective_prompt = CHAT_TEMPLATE_FMT.format(prompt=prompt) if USE_CHAT_TEMPLATE else prompt
+        for temp in TEMPERATURES:
+            with torch.inference_mode():
+                text, n_tokens, hit_eos = generate_one(
+                    model, tokenizer, effective_prompt, temp, device=device
+                )
+            rep1 = compute_ngram_rep(text, 1)
+            rep2 = compute_ngram_rep(text, 2)
+            rep3 = compute_ngram_rep(text, 3)
+            rep4 = compute_ngram_rep(text, 4)
+            diversity = compute_diversity_metrics(text)
+
+            entry = {
+                "prompt": prompt,
+                "chat_template": USE_CHAT_TEMPLATE,
+                "effective_prompt": effective_prompt if USE_CHAT_TEMPLATE else prompt,
+                "temperature": temp,
+                "generated_tokens": n_tokens,
+                "hit_eos": hit_eos,
+                "1gram_rep": round(rep1, 4),
+                "2gram_rep": round(rep2, 4),
+                "3gram_rep": round(rep3, 4),
+                "4gram_rep": round(rep4, 4),
+                "distinct_1": round(diversity["distinct_1"], 4),
+                "distinct_2": round(diversity["distinct_2"], 4),
+                "distinct_3": round(diversity["distinct_3"], 4),
+                "type_token_ratio": round(diversity["type_token_ratio"], 4),
+                "text": text[:500],  # truncate for readability
+            }
+            results.append(entry)
+            done += 1
+
+            label = "greedy" if temp == 0.0 else f"t={temp}"
+            print(
+                f"[GEN {device}] ({done}/{total_combinations}) "
+                f"{prompt[:15]}... ({label}): "
+                f"{n_tokens}tok, 3gram_rep={rep3:.2%}, eos={hit_eos}"
+            )
+
+    elapsed = time.time() - t0
+
+    # Aggregate stats per temperature group
+    greedy = [r for r in results if r["temperature"] == 0.0]
+    sampled = [r for r in results if r["temperature"] > 0.0]
+
+    if not greedy:
+        logger.warning("No greedy generation results — all prompts may have failed")
+    if not sampled:
+        logger.warning("No sampled generation results")
+
+    summary = {
+        "total_generations": len(results),
+        "n_prompts": len(PROMPTS),
+        "temperatures": TEMPERATURES,
+        "greedy_avg_1gram_rep": round(np.mean([r["1gram_rep"] for r in greedy]), 4) if greedy else 0.0,
+        "greedy_avg_2gram_rep": round(np.mean([r["2gram_rep"] for r in greedy]), 4) if greedy else 0.0,
+        "greedy_avg_3gram_rep": round(np.mean([r["3gram_rep"] for r in greedy]), 4) if greedy else 0.0,
+        "greedy_avg_4gram_rep": round(np.mean([r["4gram_rep"] for r in greedy]), 4) if greedy else 0.0,
+        "greedy_eos_rate": round(np.mean([r["hit_eos"] for r in greedy]), 4) if greedy else 0.0,
+        "greedy_avg_tokens": round(np.mean([r["generated_tokens"] for r in greedy]), 1) if greedy else 0.0,
+        "sampled_avg_3gram_rep": round(np.mean([r["3gram_rep"] for r in sampled]), 4) if sampled else 0.0,
+        "sampled_eos_rate": round(np.mean([r["hit_eos"] for r in sampled]), 4) if sampled else 0.0,
+        "sampled_avg_tokens": round(np.mean([r["generated_tokens"] for r in sampled]), 1) if sampled else 0.0,
+        "greedy_avg_distinct_1": round(float(np.mean([r["distinct_1"] for r in greedy])), 4) if greedy else 0.0,
+        "greedy_avg_distinct_2": round(float(np.mean([r["distinct_2"] for r in greedy])), 4) if greedy else 0.0,
+        "greedy_avg_distinct_3": round(float(np.mean([r["distinct_3"] for r in greedy])), 4) if greedy else 0.0,
+        "sampled_avg_distinct_2": round(float(np.mean([r["distinct_2"] for r in sampled])), 4) if sampled else 0.0,
+        "token_count_min": int(np.min([r["generated_tokens"] for r in results])) if results else 0,
+        "token_count_max": int(np.max([r["generated_tokens"] for r in results])) if results else 0,
+        "token_count_p25": int(np.percentile([r["generated_tokens"] for r in results], 25)) if results else 0,
+        "token_count_p75": int(np.percentile([r["generated_tokens"] for r in results], 75)) if results else 0,
+        "elapsed_sec": round(elapsed, 1),
+    }
+
+    print(
+        f"[GEN {device}] DONE greedy 3gram_rep={summary['greedy_avg_3gram_rep']:.4f}, "
+        f"eos_rate={summary['greedy_eos_rate']:.2%}, {elapsed:.1f}s"
+    )
+    return {"summary": summary, "samples": results}
+
+
+def eval_repetition_grid(device: str) -> dict:
+    """Grid search over 12 generation parameter combinations x 5 prompts.
+
+    Evaluates each config (temperature x repetition_penalty) on the first 5
+    prompts and returns results sorted by average 3-gram repetition rate.
+
+    Args:
+        device: CUDA device string, e.g. "cuda:5".
+
+    Returns:
+        Dict with keys:
+          - grid_results: list of per-config dicts, sorted by avg_3gram_rep
+          - best: config with lowest avg_3gram_rep
+          - elapsed_sec: wall-clock time
+    """
+    torch.cuda.set_device(int(device.split(":")[-1]))
+    print(f"[REP {device}] Loading model...")
+    model = _load_model(device)
+    tokenizer = _load_tokenizer()
+    t0 = time.time()
+
+    rep_prompts = PROMPTS[:5]  # first 5 prompts
+    results: list[dict] = []
+
+    total = len(REP_GRID) * len(rep_prompts)
+    done = 0
+
+    if USE_CHAT_TEMPLATE:
+        print(f"[REP {device}] Chat template ENABLED", flush=True)
+
+    for params in REP_GRID:
+        combo_results: list[dict] = []
+        for prompt in rep_prompts:
+            effective_prompt = CHAT_TEMPLATE_FMT.format(prompt=prompt) if USE_CHAT_TEMPLATE else prompt
+            with torch.inference_mode():
+                text, n_tokens, hit_eos = generate_one(
+                    model,
+                    tokenizer,
+                    effective_prompt,
+                    temperature=params["temperature"],
+                    repetition_penalty=params["repetition_penalty"],
+                    device=device,
+                    max_new_tokens=256,
+                )
+            combo_results.append(
+                {
+                    "prompt": prompt,
+                    "n_tokens": n_tokens,
+                    "hit_eos": hit_eos,
+                    "1gram_rep": compute_ngram_rep(text, 1),
+                    "2gram_rep": compute_ngram_rep(text, 2),
+                    "3gram_rep": compute_ngram_rep(text, 3),
+                    "4gram_rep": compute_ngram_rep(text, 4),
+                }
+            )
+            done += 1
+
+        if not combo_results:
+            logger.warning("All prompts failed for config %s — skipping", params.get("name", "unknown"))
+            continue
+
+        avg_3gram = float(np.mean([r["3gram_rep"] for r in combo_results]))
+        avg_4gram = float(np.mean([r["4gram_rep"] for r in combo_results]))
+        eos_rate = float(np.mean([r["hit_eos"] for r in combo_results]))
+        avg_tokens = float(np.mean([r["n_tokens"] for r in combo_results]))
+
+        entry = {
+            "params": params["name"],
+            "temperature": params["temperature"],
+            "repetition_penalty": params["repetition_penalty"],
+            "avg_3gram_rep": round(avg_3gram, 4),
+            "avg_4gram_rep": round(avg_4gram, 4),
+            "eos_rate": round(eos_rate, 4),
+            "avg_tokens": round(avg_tokens, 1),
+            "per_prompt": combo_results,
+        }
+        results.append(entry)
+        print(
+            f"[REP {device}] {params['name']}: "
+            f"3gram={avg_3gram:.2%}, 4gram={avg_4gram:.2%}, "
+            f"eos={eos_rate:.0%}, {avg_tokens:.0f}tok"
+        )
+
+    elapsed = time.time() - t0
+
+    # Sort by avg 3-gram repetition (ascending = better)
+    sorted_results = sorted(results, key=lambda r: r["avg_3gram_rep"])
+    best = sorted_results[0]
+
+    print(
+        f"[REP {device}] DONE best={best['params']} "
+        f"(3gram={best['avg_3gram_rep']:.2%}), {elapsed:.1f}s"
+    )
+    return {
+        "grid_results": sorted_results,
+        "best": {
+            "params": best["params"],
+            "temperature": best["temperature"],
+            "repetition_penalty": best["repetition_penalty"],
+            "avg_3gram_rep": best["avg_3gram_rep"],
+            "avg_4gram_rep": best["avg_4gram_rep"],
+        },
+        "elapsed_sec": round(elapsed, 1),
+    }