56m/SmallMultiplicationLeaderboard

1# ==============================================================================2# Multiplication SLM Benchmark Suite (Interactive TUI Wizard)3# Dynamic Dataset Generator corresponding to 'cot_partial', 'cot_direct', 'cot_cot'4# ==============================================================================5 6import json7import math8import os9import random10import subprocess11import sys12import torch13import torch.nn.functional as F14from torch.utils.data import DataLoader, Dataset15 16# ------------------------------------------------------------------------------17# 1. Install & Import Rich18# ------------------------------------------------------------------------------19try:20    from rich.align import Align21    from rich import box22    from rich.console import Console23    from rich.panel import Panel24    from rich.progress import (25        BarColumn,26        Progress,27        SpinnerColumn,28        TextColumn,29        TimeRemainingColumn,30    )31    from rich.prompt import Prompt32    from rich.table import Table33except ImportError:34    subprocess.check_call(35        [sys.executable, "-m", "pip", "install", "rich", "-q"]36    )37    from rich.align import Align38    from rich import box39    from rich.console import Console40    from rich.panel import Panel41    from rich.progress import (42        BarColumn,43        Progress,44        SpinnerColumn,45        TextColumn,46        TimeRemainingColumn,47    )48    from rich.prompt import Prompt49    from rich.table import Table50 51from transformers import GPT2LMHeadModel52 53console = Console()54 55# ------------------------------------------------------------------------------56# 2. Tokenizer Definition57# ------------------------------------------------------------------------------58DEFAULT_CHARS = sorted(59    list(60        set(61            "0123456789 *+=:\nQStepsAnsLet'sthinkstepby"62            + "<|endoftext|>"63            + "<|pad|>"64        )65    )66)67 68class CharTokenizer:69    def __init__(self, chars=None):70        if chars is None:71            chars = DEFAULT_CHARS72        self.chars = sorted(list(set(chars)))73        if "<|pad|>" not in self.chars:74            self.chars.append("<|pad|>")75        self.c2i = {ch: i for i, ch in enumerate(self.chars)}76        self.i2c = {i: ch for i, ch in enumerate(self.chars)}77        self.pad_id = self.c2i["<|pad|>"]78        self.vocab_size = len(self.chars)79 80    @classmethod81    def load(cls, filepath):82        with open(filepath, "r", encoding="utf-8") as f:83            data = json.load(f)84        return cls(data["chars"])85 86    def encode(self, text):87        return [self.c2i[c] for c in text if c in self.c2i]88 89    def decode(self, ids):90        return "".join(91            [self.i2c[i] for i in ids if i in self.i2c and i != self.pad_id]92        )93 94# ------------------------------------------------------------------------------95# 3. Interactive TUI Wizard96# ------------------------------------------------------------------------------97console.clear()98console.print()99console.print(100    Panel(101        Align.center(102            "[bold yellow]🤗 Hugging Face Micro-SLM Benchmark Harness[/bold yellow]\n"103            "[dim]10,000 Qs Log-Likelihood & Exact Match  |  500 Qs Perplexity (PPL)\n"104            "Rule: Strictly for Small Models (≤ 250M Parameters)[/dim]"105        ),106        box=box.ROUNDED,107        border_style="yellow",108        padding=(1, 2),109    )110)111 112# Step 1: Model Directory113candidate_paths = [114    d115    for d in ["./slm_mult_model", "./slm_mult_checkpoints", "."]116    if os.path.exists(os.path.join(d, "config.json"))117    or os.path.exists(os.path.join(d, "tokenizer.json"))118]119detected_hint = (120    f"[green]Found model at:[/green] [bold]{candidate_paths[0]}[/bold]"121    if candidate_paths122    else "[dim]No model detected in current folder.[/dim]"123)124 125input_dialog = Table.grid(padding=1)126input_dialog.add_column(style="bold cyan")127input_dialog.add_column()128input_dialog.add_row("Auto Detected:", detected_hint)129input_dialog.add_row("Default Path:", "[bold yellow]./slm_mult_model[/bold yellow]")130 131console.print(132    Panel(133        input_dialog,134        title="[bold white] 📁 Step 1/4: Model Directory [/bold white]",135        title_align="left",136        border_style="cyan",137        box=box.ROUNDED,138    )139)140 141while True:142    model_dir = Prompt.ask(143        " [bold cyan]➤ Model Directory Path[/bold cyan]",144        default=candidate_paths[0] if candidate_paths else "./slm_mult_model",145    )146    if os.path.exists(model_dir):147        break148    console.print("[bold red]Directory not found! Please try again.[/bold red]")149 150# Step 2: Model Display Name151default_name = f"YourOrg/{os.path.basename(os.path.abspath(model_dir))}"152console.print()153console.print(154    Panel(155        "[dim]Enter the model name for the leaderboard (e.g. OrgName/ModelName).[/dim]",156        title="[bold white] 🏷️ Step 2/4: Model ID / Display Name [/bold white]",157        title_align="left",158        border_style="cyan",159        box=box.ROUNDED,160    )161)162model_name = Prompt.ask(163    " [bold cyan]➤ Model Name[/bold cyan]",164    default=default_name165)166 167# Step 3: CoT Mode Selection168cot_table = Table(show_header=False, box=None, padding=(0, 1))169cot_table.add_row("[bold yellow]1. cot_partial[/bold yellow]", ": Step-by-step arithmetic decomposition (筆算分解)")170cot_table.add_row("[bold yellow]2. cot_direct [/bold yellow]", ": Direct answer output without intermediate steps (直接出力)")171cot_table.add_row("[bold yellow]3. cot_cot    [/bold yellow]", ": Chain-of-Thought prompt -> Final Answer only (プロンプトCoT/最終回答のみ)")172 173console.print()174console.print(175    Panel(176        cot_table,177        title="[bold white] 🧠 Step 3/4: CoT (Reasoning) Mode [/bold white]",178        title_align="left",179        border_style="cyan",180        box=box.ROUNDED,181    )182)183 184cot_choice = Prompt.ask(185    " [bold cyan]➤ Select CoT Mode[/bold cyan]",186    choices=["cot_partial", "cot_direct", "cot_cot"],187    default="cot_partial"188)189 190# Step 4: Highlight Tag191console.print()192console.print(193    Panel(194        "[dim]Highlight this entry as a BASE model on the leaderboard?[/dim]",195        title="[bold white] ⭐ Step 4/4: Base Model Highlight [/bold white]",196        title_align="left",197        border_style="cyan",198        box=box.ROUNDED,199    )200)201is_highlight_str = Prompt.ask(202    " [bold cyan]➤ Highlight as Base Model?[/bold cyan]",203    choices=["y", "n"],204    default="n"205)206is_highlight = (is_highlight_str.lower() == "y")207 208# ------------------------------------------------------------------------------209# 4. Model Loading & Verification210# ------------------------------------------------------------------------------211tokenizer_path = os.path.join(model_dir, "tokenizer.json")212if os.path.exists(tokenizer_path):213    tokenizer = CharTokenizer.load(tokenizer_path)214    tokenizer_status = f"[green]Loaded from ({tokenizer_path})[/green]"215else:216    tokenizer = CharTokenizer(DEFAULT_CHARS)217    tokenizer_status = "[yellow]Standard Vocab (Fallback)[/yellow]"218 219device = "cuda" if torch.cuda.is_available() else "cpu"220 221with console.status(f"[bold cyan]Loading weights onto [bold green]{device.upper()}[/bold green]...[/bold cyan]"):222    model = GPT2LMHeadModel.from_pretrained(model_dir).to(device)223    model.eval()224 225total_params = sum(p.numel() for p in model.parameters())226params_str = (227    f"{total_params / 1e6:.1f}M"228    if total_params >= 1e6229    else f"{total_params / 1e3:.1f}K"230)231is_qualified = total_params <= 250_000_000232 233info_table = Table(show_header=False, box=None, padding=(0, 2))234info_table.add_row("Model Name:", f"[bold white]{model_name}[/bold white]")235info_table.add_row("Directory:", f"[dim]{os.path.abspath(model_dir)}[/dim]")236info_table.add_row("CoT Mode:", f"[bold yellow]{cot_choice}[/bold yellow]")237info_table.add_row("Highlighted:", f"[bold magenta]{is_highlight}[/bold magenta]")238info_table.add_row("Total Parameters:", f"[bold cyan]{params_str}[/bold cyan] ({total_params:,})")239info_table.add_row(240    "Eligibility (≤ 250M):",241    "[bold green]✔ QUALIFIED[/bold green]" if is_qualified else "[bold red]✖ DISQUALIFIED (>250M)[/bold red]"242)243 244console.print()245console.print(246    Panel(247        info_table,248        title="[bold green] 🤖 Model Specification Confirmed [/bold green]",249        title_align="left",250        border_style="green",251        box=box.ROUNDED,252    )253)254 255# ------------------------------------------------------------------------------256# 5. Dynamic Dataset Generator (Based on cot_choice)257# ------------------------------------------------------------------------------258def generate_benchmark_sample(cot_mode="cot_partial"):259    mode = random.choices(["2terms", "3terms"], weights=[0.7, 0.3])[0]260 261    if mode == "2terms":262        a = random.randint(2, 999)263        b = random.randint(2, 99)264        expr = f"{a} * {b}"265        ans = a * b266 267        if cot_mode == "cot_direct":268            prompt = f"Q: {expr}\nAns:"269            completion = f" {ans}<|endoftext|>"270 271        elif cot_mode == "cot_cot":272            prompt = f"Q: {expr}\nLet's think step by step.\nAns:"273            completion = f" {ans}<|endoftext|>"274 275        else:  # cot_partial (standard step-by-step)276            if b >= 10:277                steps = []278                sub_products = []279                for i, digit in enumerate(reversed(str(b))):280                    p_val = int(digit) * (10**i)281                    if p_val > 0:282                        part_ans = a * p_val283                        sub_products.append(part_ans)284                        steps.append(f"Step {len(steps)+1}: {a} * {p_val} = {part_ans}")285                if len(sub_products) > 1:286                    add_expr = " + ".join(map(str, sub_products))287                    steps.append(f"Step {len(steps)+1}: {add_expr} = {ans}")288                steps_text = "\n".join(steps)289            else:290                steps_text = f"Step 1: {expr} = {ans}"291 292            prompt = f"Q: {expr}\nSteps:\n"293            completion = f"{steps_text}\nAns: {ans}<|endoftext|>"294 295    else:  # 3terms296        a = random.randint(2, 50)297        b = random.randint(2, 20)298        c = random.randint(2, 20)299        expr = f"{a} * {b} * {c}"300        r1, r2 = a * b, a * b * c301 302        if cot_mode == "cot_direct":303            prompt = f"Q: {expr}\nAns:"304            completion = f" {r2}<|endoftext|>"305 306        elif cot_mode == "cot_cot":307            prompt = f"Q: {expr}\nLet's think step by step.\nAns:"308            completion = f" {r2}<|endoftext|>"309 310        else:  # cot_partial311            prompt = f"Q: {expr}\nSteps:\n"312            completion = f"Step 1: {a} * {b} = {r1}\nStep 2: {r1} * {c} = {r2}\nAns: {r2}<|endoftext|>"313 314    return prompt, completion315 316class EvalDataset(Dataset):317    def __init__(self, samples, max_len=160):318        self.items = []319        for prompt, comp in samples:320            full_text = prompt + comp321            p_tokens = tokenizer.encode(prompt)322            full_tokens = tokenizer.encode(full_text)323            prompt_len = len(p_tokens)324            seq_len = len(full_tokens)325 326            if seq_len > max_len:327                full_tokens = full_tokens[:max_len]328 329            labels = list(full_tokens)330            for i in range(min(prompt_len, len(labels))):331                labels[i] = -100332 333            pad_len = max_len - len(full_tokens)334            input_ids = full_tokens + [tokenizer.pad_id] * pad_len335            labels = labels + [-100] * pad_len336 337            self.items.append(338                {339                    "input_ids": torch.tensor(input_ids, dtype=torch.long),340                    "labels": torch.tensor(labels, dtype=torch.long),341                }342            )343 344    def __len__(self):345        return len(self.items)346 347    def __getitem__(self, idx):348        return self.items[idx]349 350with console.status(f"[bold blue]Generating 10,500 questions in '[yellow]{cot_choice}[/yellow]' mode...[/bold blue]"):351    random.seed(2025)352    eval_10k_data = [generate_benchmark_sample(cot_choice) for _ in range(10000)]353    eval_ppl_data = [generate_benchmark_sample(cot_choice) for _ in range(500)]354    dataset_10k = EvalDataset(eval_10k_data)355    dataset_ppl = EvalDataset(eval_ppl_data)356 357# Auto Batch Size Discovery358def find_optimal_batch_size(model, min_b=1, max_b=1024, seq_len=160, device="cuda"):359    if device == "cpu": return 64360    candidate_sizes = [2**i for i in range(int(math.log2(min_b)), int(math.log2(max_b)) + 1)]361    optimal_b = min_b362    with console.status("[bold magenta]Probing maximum safe batch size...[/bold magenta]"):363        for b in candidate_sizes:364            try:365                torch.cuda.empty_cache()366                dummy_input = torch.randint(0, tokenizer.vocab_size, (b, seq_len), device=device)367                with torch.no_grad():368                    _ = model(dummy_input)369                optimal_b = b370            except Exception:371                break372    return optimal_b373 374auto_batch_size = find_optimal_batch_size(model, device=device)375console.print(f"\n[bold green]✔ Hardware Optimized:[/bold green] Batch Size set to [bold cyan]{auto_batch_size}[/bold cyan]\n")376 377# ------------------------------------------------------------------------------378# 6. Benchmark Execution Loop379# ------------------------------------------------------------------------------380loader_10k = DataLoader(dataset_10k, batch_size=auto_batch_size, shuffle=False)381loader_ppl = DataLoader(dataset_ppl, batch_size=auto_batch_size, shuffle=False)382 383total_log_likelihood = 0.0384total_tokens = 0385correct_tokens = 0386perfect_sequences = 0387total_sequences = 0388 389with Progress(390    SpinnerColumn(),391    TextColumn("[progress.description]{task.description}"),392    BarColumn(complete_style="cyan", finished_style="green"),393    TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),394    TimeRemainingColumn(),395    console=console,396) as progress:397 398    task1 = progress.add_task(f"[cyan]Evaluating 10,000 Qs ({cot_choice})...", total=len(loader_10k))399    with torch.no_grad():400        for batch in loader_10k:401            input_ids = batch["input_ids"].to(device)402            labels = batch["labels"].to(device)403            outputs = model(input_ids)404            logits = outputs.logits405 406            shift_logits = logits[:, :-1, :].contiguous()407            shift_labels = labels[:, 1:].contiguous()408 409            log_probs = F.log_softmax(shift_logits, dim=-1)410            mask = shift_labels != -100411 412            gathered = torch.gather(log_probs, 2, shift_labels.unsqueeze(-1).clamp(min=0)).squeeze(-1)413            gathered = gathered * mask414 415            total_log_likelihood += gathered.sum().item()416            total_tokens += mask.sum().item()417 418            preds = torch.argmax(shift_logits, dim=-1)419            matches = (preds == shift_labels) & mask420            correct_tokens += matches.sum().item()421 422            seq_matches = (matches.sum(dim=1) == mask.sum(dim=1)).sum().item()423            perfect_sequences += seq_matches424            total_sequences += input_ids.size(0)425 426            progress.update(task1, advance=1)427 428    task2 = progress.add_task("[magenta]Calculating Perplexity on 500 Qs...", total=len(loader_ppl))429    total_loss = 0.0430    ppl_batches = 0431    loss_fn = torch.nn.CrossEntropyLoss(ignore_index=-100)432 433    with torch.no_grad():434        for batch in loader_ppl:435            input_ids = batch["input_ids"].to(device)436            labels = batch["labels"].to(device)437            outputs = model(input_ids)438            shift_logits = outputs.logits[:, :-1, :].contiguous()439            shift_labels = labels[:, 1:].contiguous()440 441            loss = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))442            total_loss += loss.item()443            ppl_batches += 1444            progress.update(task2, advance=1)445 446# Metrics Calculation447token_accuracy = (correct_tokens / total_tokens) * 100448sequence_accuracy = (perfect_sequences / total_sequences) * 100449avg_loss = total_loss / ppl_batches450ppl_score = math.exp(avg_loss)451 452p = sequence_accuracy / 100.0453p_smooth = min(max(p, 1e-4), 1.0 - 1e-5)454logit_val = math.log(p_smooth / (1.0 - p_smooth))455rating_score = max(0.0, 1050.0 + (252.47 * logit_val))456 457# ------------------------------------------------------------------------------458# 7. HTML Leaderboard Ready JSON Output459# ------------------------------------------------------------------------------460hf_leaderboard_entry = {461    "name": model_name,462    "params_str": params_str,463    "params_num": total_params,464    "cot": cot_choice,465    "score": round(rating_score, 1),466    "exact": round(sequence_accuracy, 1),467    "token_acc": round(token_accuracy, 1),468    "ppl": round(ppl_score, 3),469    "highlight": is_highlight470}471 472json_str = json.dumps(hf_leaderboard_entry, indent=4)473 474console.print()475console.print(476    Panel(477        f"[bold green]{json_str}[/bold green]",478        title="[bold yellow] 📋 HTML Ready Leaderboard Entry (Copy & Paste to modelData) [/bold yellow]",479        title_align="left",480        border_style="yellow",481        box=box.ROUNDED,482    )483)484console.print("[dim]Copy the JSON above and paste it directly into the 'modelData' array inside index.html![/dim]\n")