Part B: Sampling-based improvement

"""Part B: exactly how majority voting changes accuracy, with the binomial formula and a simulation."""
from math import comb

import numpy as np

rng = np.random.default_rng(5)


def majority_exact(p: float, n: int) -> float:
    """P(more than half of n independent samples are correct); ties split 50/50."""
    win = sum(comb(n, k) * p**k * (1 - p) ** (n - k) for k in range(n // 2 + 1, n + 1))
    if n % 2 == 0:
        win += 0.5 * comb(n, n // 2) * (p * (1 - p)) ** (n // 2)
    return win


def majority_simulated(p: float, n: int, questions: int = 200_000) -> float:
    correct = rng.random((questions, n)) < p
    k = correct.sum(axis=1)
    return float(np.mean((k > n / 2) + 0.5 * (k == n / 2)))


# 1. Independent errors: the formula and the simulation agree.
ns = (1, 3, 5, 9, 15, 31)
print("Independent samples, two outcomes (right or wrong). Rows: per-sample accuracy p.")
print("p     " + "".join(f"n={n:<6}" for n in ns))
for p in (0.3, 0.5, 0.6, 0.7, 0.8, 0.9):
    print(f"{p:<5} " + "".join(f"{majority_exact(p, n):<8.3f}" for n in ns))
print("simulated p=0.7, n=5:", round(majority_simulated(0.7, 5), 3), " exact:", round(majority_exact(0.7, 5), 3))


# 2. Correlated errors: each question has its own accuracy p_i drawn from a Beta with mean p.
#    rho is the correlation between two samples' correctness on the same question.
def majority_correlated(p: float, n: int, rho: float, questions: int = 200_000) -> float:
    if rho == 0:
        return majority_exact(p, n)
    a, b = p * (1 / rho - 1), (1 - p) * (1 / rho - 1)
    p_i = rng.beta(a, b, size=questions)
    k = rng.binomial(n, p_i)
    return float(np.mean((k > n / 2) + 0.5 * (k == n / 2)))


print("\nCorrelated samples, p=0.7. Rows: correlation rho between samples on the same question.")
print("rho   " + "".join(f"n={n:<6}" for n in ns))
for rho in (0.0, 0.1, 0.3, 0.6, 0.9):
    print(f"{rho:<5} " + "".join(f"{majority_correlated(0.7, n, rho):<8.3f}" for n in ns))
for rho in (0.1, 0.3, 0.6):
    a, b = 0.7 * (1 / rho - 1), 0.3 * (1 / rho - 1)
    ceiling = float(np.mean(rng.beta(a, b, size=200_000) > 0.5))
    print(f"ceiling as n grows, rho={rho}: {ceiling:.3f} (share of questions whose own p_i > 0.5)")


# 3. More than two answers: plurality voting when wrong answers spread out or pile up.
def plurality(probs: list[float], n: int, questions: int = 100_000) -> float:
    """Answer 0 is correct. Ties are broken uniformly at random among the tied answers."""
    counts = rng.multinomial(n, probs, size=questions)
    top = counts.max(axis=1, keepdims=True)
    tied = counts == top
    return float(np.mean(tied[:, 0] / tied.sum(axis=1)))


print("\nPlurality vote with four possible answers (correct answer has probability 0.40):")
for name, probs in (("wrong answers spread 0.20/0.20/0.20", [0.4, 0.2, 0.2, 0.2]),
                    ("one tempting wrong answer 0.45/0.10/0.05", [0.4, 0.45, 0.1, 0.05])):
    print(f"  {name:<42}" + "  ".join(f"n={n}: {plurality(probs, n):.3f}" for n in (1, 5, 15, 31)))

"""Part B: self-consistency for real on TinyLM. Sample several completions, vote on the number they state."""
import re
import time
from collections import Counter

import torch

from supportdesk.tinylm import SamplingParams, generate, load

torch.set_num_threads(1)  # this machine is shared; 1 thread keeps timings steadier
model, tokenizer = load()
NUMBER = re.compile(r"\d+(?:\.\d+)?")

# Prompt, and the number a correct continuation states (from the help-center facts).
FACT_PROMPTS = [
    ("Reset links expire after", "30"),
    ("After 5 failed attempts an account is locked for", "15"),
    ("Annual plans cancelled within", "14"),
    ("Business costs 24 USD per user per month, or", "20"),
    ("Free supports up to", "3"),
    ("Customer (Ana): Hi, how long until my locked account opens?\nAgent (Ben):", "15"),
    ("Customer (Ana): Hi, what does Business cost per user if we pay annually?\nAgent (Ben):", "20"),
    ("Customer (Ana): Hi, how many users can the Free plan have?\nAgent (Ben):", "3"),
]
N_SAMPLES = 9


def first_number(text: str) -> str | None:
    match = NUMBER.search(text)
    return match.group(0) if match else None


started = time.perf_counter()
for temperature in (1.0, 1.5):
    greedy_ok = single_ok = vote_ok = 0
    print(f"temperature {temperature}, {N_SAMPLES} samples per prompt")
    for prompt, answer in FACT_PROMPTS:
        greedy = first_number(generate(model, tokenizer, prompt, SamplingParams(max_new_tokens=16, temperature=0)).text)
        samples = [first_number(generate(model, tokenizer, prompt,
                                         SamplingParams(max_new_tokens=16, temperature=temperature, seed=s)).text)
                   for s in range(N_SAMPLES)]
        votes = Counter(x for x in samples if x is not None)
        top = votes.most_common(2)
        tie = len(top) == 2 and top[0][1] == top[1][1]
        winner = None if not top or tie else top[0][0]  # a tie is not a majority: count it as no answer
        p_hat = sum(x == answer for x in samples) / N_SAMPLES
        greedy_ok += greedy == answer
        single_ok += p_hat
        vote_ok += winner == answer
        label = prompt.splitlines()[0].replace("Customer (Ana): Hi, ", "Q: ")
        print(f"  want {answer:>3}  greedy {str(greedy):>4}  per-sample {p_hat:.2f}  vote {'tie' if tie else str(winner):>4}  "
              f"{dict(votes.most_common(3))}  | {label[:48]}")
    k = len(FACT_PROMPTS)
    print(f"  accuracy over {k} prompts: greedy {greedy_ok}/{k}, one sample {single_ok:.2f}/{k} (average), "
          f"majority of {N_SAMPLES} {vote_ok}/{k}\n")
print(f"{2 * len(FACT_PROMPTS) * (N_SAMPLES + 1)} generations in {time.perf_counter() - started:.0f} s")

"""Part B: majority vote vs best-of-n with a verifier vs best-of-n with a noisy scorer, on the 20 refund cases.

The generator is SIMULATED: each sample is right with probability 0.9 on easy
cases and 0.45 on trap cases (cases the one-step shortcut gets wrong), and trap
mistakes usually repeat the shortcut's tempting answer. Only the policy, the
cases, and the verifier are real. Replace the simulator with real samples
(temperature above 0) to measure your model.
"""
import numpy as np

from examples.m05_helpers import CASES, LABELS, decide, shortcut_baseline

rng = np.random.default_rng(11)
REPEATS = 2000


def sample_label(case) -> str:
    trap = shortcut_baseline(case.text) != case.gold
    if rng.random() < (0.45 if trap else 0.9):
        return case.gold
    wrong = [label for label in LABELS if label != case.gold]
    tempting = shortcut_baseline(case.text)
    if trap and rng.random() < 0.8:
        return tempting
    return str(rng.choice(wrong))


def verifier(case, label: str) -> bool:
    """External check: recompute the decision from the billing record (the case's facts)."""
    return label == decide(case.facts)


def noisy_score(case, label: str, sigma: float = 0.8) -> float:
    """A learned scorer: prefers right answers on average, but with noise."""
    return float(label == case.gold) + rng.normal(0, sigma)


def run(n: int, repeats: int = REPEATS) -> dict[str, float]:
    totals = dict(majority=0, verified=0, escalated=0, scored=0, coverage=0, calls_verified=0)
    for _ in range(repeats):
        for case in CASES:
            samples = [sample_label(case) for _ in range(n)]
            counts = {label: samples.count(label) for label in set(samples)}
            top = max(counts.values())
            winner = str(rng.choice(sorted(label for label, c in counts.items() if c == top)))
            totals["majority"] += winner == case.gold
            accepted = next((s for s in samples if verifier(case, s)), None)
            totals["verified"] += accepted == case.gold
            totals["escalated"] += accepted is None
            # Sampling one at a time, you can stop at the first sample the verifier accepts.
            totals["calls_verified"] += samples.index(accepted) + 1 if accepted is not None else n
            totals["scored"] += max(samples, key=lambda s: noisy_score(case, s)) == case.gold
            totals["coverage"] += case.gold in samples
    return {k: v / (repeats * len(CASES)) for k, v in totals.items()}


if __name__ == "__main__":
    print("Share of the 20 cases answered correctly (simulated generator, 2,000 repeats per cell)")
    print(f"{'n':>3} {'majority':>9} {'best-of-n+verifier':>19} {'avg calls':>9} {'escalated':>10} {'best-of-n+scorer':>17} {'pass@n':>7}")
    for n in (1, 3, 5, 9, 15):
        r = run(n)
        print(f"{n:>3} {r['majority']:>9.3f} {r['verified']:>19.3f} {r['calls_verified']:>9.2f} {r['escalated']:>10.3f} "
              f"{r['scored']:>17.3f} {r['coverage']:>7.3f}")

Share of the 20 cases answered correctly (simulated generator, 2,000 repeats per cell)
  n  majority  best-of-n+verifier avg calls  escalated  best-of-n+scorer  pass@n
  1     0.695               0.695      1.00      0.305             0.695   0.695
  3     0.742               0.922      1.45      0.078             0.841   0.922
  5     0.758               0.978      1.56      0.022             0.887   0.978
  9     0.761               0.997      1.61      0.003             0.923   0.997
 15     0.756               1.000      1.61      0.000             0.945   1.000

"""Part B: ensembling three different triage classifiers, and why their error overlap decides the gain.

All three are real, deterministic classifiers (no LLM), so every number here is
a real measurement on the 72 labelled tickets.
"""
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold, cross_val_predict
from sklearn.pipeline import make_pipeline

from examples.m05_helpers import TRIAGE_KEYWORDS, keyword_reader, triage_prompt, wilson
from supportdesk.data import load_tickets
from supportdesk.kb_search import KBSearch

tickets = load_tickets()
gold = [t.gold["category"] for t in tickets]

# Member 1: keyword definitions read from a triage prompt.
prompt = triage_prompt(TRIAGE_KEYWORDS)
keywords = [keyword_reader(prompt, t.text) for t in tickets]

# Member 2: retrieve the best help-center article and map it to a category (hand-written map).
ARTICLE_TO_CATEGORY = {
    "account-login": "account_access", "account-sso": "account_access", "data-privacy": "account_access",
    "billing-invoices": "billing", "billing-plans": "billing", "billing-refunds": "cancellation",
    "boards-automations": "how_to", "exports-data": "how_to", "mobile-app": "how_to",
    "integrations-slack": "bug", "status-incidents": "bug", "feature-requests": "feature_request",
}
search = KBSearch()
retrieval = []
for t in tickets:
    hits = search.search(t.text, k=1)
    retrieval.append(ARTICLE_TO_CATEGORY[hits[0].article_id] if hits else "how_to")

# Member 3: character n-gram TF-IDF + logistic regression, scored with 6-fold cross-validation.
model = make_pipeline(TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 4)), LogisticRegression(max_iter=2000, C=10))
folds = StratifiedKFold(n_splits=6, shuffle=True, random_state=0)
learned = list(cross_val_predict(model, [t.text for t in tickets], gold, cv=folds))


def vote(a: str, b: str, c: str) -> str:
    """Plurality of three; if all three disagree, trust the learned model."""
    if a == b or a == c:
        return a
    if b == c:
        return b
    return c


members = {"keywords": keywords, "retrieval": retrieval, "learned": learned}
members["ensemble"] = [vote(*x) for x in zip(keywords, retrieval, learned)]
n = len(tickets)
for name, preds in members.items():
    k = sum(p == g for p, g in zip(preds, gold))
    lo, hi = wilson(k, n)
    print(f"{name:<10} {k:>2}/{n} = {k / n:.3f}  95% CI {lo:.2f}-{hi:.2f}")
any_right = sum(g in (a, b, c) for g, a, b, c in zip(gold, keywords, retrieval, learned))
print(f"at least one member right: {any_right}/{n} (the ceiling for any way of combining them)")

print("\nError overlap (tickets both members got wrong) and correlation of their error indicators:")
errors = {name: np.array([p != g for p, g in zip(members[name], gold)]) for name in ("keywords", "retrieval", "learned")}
names = list(errors)
for i in range(3):
    for j in range(i + 1, 3):
        a, b = errors[names[i]], errors[names[j]]
        both = int((a & b).sum())
        corr = float(np.corrcoef(a, b)[0, 1])
        print(f"  {names[i]:<9} & {names[j]:<9} wrong together on {both:>2}  (errors {a.sum()} and {b.sum()}), correlation {corr:+.2f}")

"""Part B: test-time compute as a dial. What each extra sample or reasoning level costs per 1,000 tickets.

Offline it combines measured prompt sizes and real prices with the SIMULATED
accuracies from m05_best_of_n.py. With --live it measures a reasoning model at
three effort levels on the 20 refund cases:
    LLM_PROVIDER=groq GROQ_API_KEY=... PYTHONPATH=. python examples/m05_effort_dial.py --live
"""
import sys

from examples.m05_best_of_n import run
from examples.m05_helpers import CASES, evaluate, prompt_direct, prompt_scaffold
from supportdesk.llm import Usage
from supportdesk.pricing import PRICES, cost_usd
from supportdesk.tokens import count_messages

MODEL = "openai/gpt-oss-120b"
tokens_in = round(sum(count_messages(prompt_scaffold(c)) for c in CASES) / len(CASES))
OUT_PER_SAMPLE = 60  # assumption: the scaffold's six short lines; replace with your measured usage
print(f"scaffold prompt: {tokens_in} input tokens on average (measured, o200k_base); "
      f"{OUT_PER_SAMPLE} output tokens per sample (assumed)\n")

print(f"{'strategy':<24} {'calls':>5} {'USD per 1k':>11} {'accuracy':>9} {'USD per 1k correct':>19}")
for n in (1, 3, 5, 9, 15):
    r = run(n, repeats=400)
    for name, calls, acc in ((f"majority n={n}", n, r["majority"]),
                             (f"verify, stop early n<={n}", r["calls_verified"], r["verified"])):
        usage = Usage(input_tokens=round(calls * tokens_in), output_tokens=round(calls * OUT_PER_SAMPLE))
        per_1k = 1000 * cost_usd(usage, MODEL)
        print(f"{name:<24} {calls:>5.2f} {per_1k:>11.3f} {acc:>9.3f} {per_1k / acc:>19.3f}")

# Prompt caching: n samples of one ticket share the whole prompt. Groq's prompt-caching page lists
# 50 percent off cached input for gpt-oss; pricing.py (checked 21 Sep 2026) lists cached_input equal to
# input for these models, so this module applies Groq's documented rate locally instead of editing it.
GROQ_CACHED_RATE = PRICES[MODEL].input * 0.5
print("\nMajority of n when samples 2..n hit the prompt cache (Groq's documented 50 percent cached rate):")
for n in (1, 5, 15):
    uncached = Usage(input_tokens=tokens_in, output_tokens=n * OUT_PER_SAMPLE)
    cached_in = (n - 1) * tokens_in
    per_1k = 1000 * (cost_usd(uncached, MODEL) + cached_in * GROQ_CACHED_RATE / 1_000_000)
    full = 1000 * cost_usd(Usage(input_tokens=n * tokens_in, output_tokens=n * OUT_PER_SAMPLE), MODEL)
    print(f"  n={n:<3} USD per 1k tickets {per_1k:.3f} with cache hits, {full:.3f} without")

print("\nReasoning effort on one call (reasoning tokens are ASSUMED; they bill as output tokens):")
for effort, reasoning in (("low", 300), ("medium", 1200), ("high", 4000)):
    usage = Usage(input_tokens=tokens_in, output_tokens=OUT_PER_SAMPLE + reasoning, reasoning_tokens=reasoning)
    print(f"  {effort:<7} ~{reasoning:>5} reasoning tokens  USD per 1k tickets {1000 * cost_usd(usage, MODEL):.3f}")

if "--live" in sys.argv:
    from supportdesk.llm import chat
    for effort in ("low", "medium", "high"):
        result = evaluate(chat, prompt_direct, reasoning_effort=effort, max_tokens=8000)
        usage = result.usage
        print(result.line(f"effort={effort}"), f" reasoning {usage.reasoning_tokens}",
              f" USD {cost_usd(usage, MODEL):.4f}")