473 lines
18 KiB
Python
473 lines
18 KiB
Python
|
"""
|
||
|
|
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),
|
||
|
|
}
|