Capability Frontiers › Code Implementation

from collections import Counter

import numpy as np

# Simulate in-context learning with a simple k-NN model
# The "context" is the set of examples provided before the test query


def knn_predict(examples, query, k=3):
    """
    Classify query using k nearest neighbors from examples.
    Mimics in-context learning: more examples = richer context.
    """
    if not examples:
        # No examples: return random guess
        return np.random.choice(["positive", "negative"])

    X = np.array([ex["embedding"] for ex in examples])
    y = [ex["label"] for ex in examples]

    # Compute distances
    distances = np.linalg.norm(X - query, axis=1)
    sorted_idx = np.argsort(distances)

    # Get k nearest neighbors
    k_actual = min(k, len(examples))
    neighbors = [y[i] for i in sorted_idx[:k_actual]]

    # Majority vote
    vote = Counter(neighbors).most_common(1)[0][0]
    return vote


np.random.seed(42)

# Create a simple 2D classification problem
# Positive class: centered at (1, 1)
# Negative class: centered at (-1, -1)
n_per_class = 50

pos_examples = [
    {
        "embedding": np.array([1.0, 1.0]) + np.random.randn(2) * 0.5,
        "label": "positive",
    }
    for _ in range(n_per_class)
]

neg_examples = [
    {
        "embedding": np.array([-1.0, -1.0]) + np.random.randn(2) * 0.5,
        "label": "negative",
    }
    for _ in range(n_per_class)
]

all_examples = pos_examples + neg_examples

# Test queries
test_queries = [
    np.array([1.2, 0.9]) + np.random.randn(2) * 0.3 for _ in range(20)
]
test_labels = ["positive"] * 20

# Evaluate accuracy as a function of number of in-context examples
shot_counts = [0, 1, 2, 4, 8, 16, 32, 64]
accuracies = []

for n_shots in shot_counts:
    ctx = all_examples[:n_shots]
    correct = sum(
        knn_predict(ctx, q) == lbl for q, lbl in zip(test_queries, test_labels)
    )
    accuracies.append(correct / len(test_labels))

def solve_direct(problem_steps):
    """
    Direct answer: try to compute the final answer in one step.
    Simulates a model without chain-of-thought.
    Introduces noise proportional to problem complexity.
    """
    # Each step adds complexity and potential error
    true_answer = sum(step["delta"] for step in problem_steps)
    noise_std = len(problem_steps) * 2.0  # more steps = more error
    answer = true_answer + np.random.randn() * noise_std
    return round(answer)


def solve_chain_of_thought(problem_steps):
    """
    Step-by-step answer: compute each step explicitly.
    Simulates chain-of-thought: each step is independently (noisily) computed.
    """
    running_total = 0
    for step in problem_steps:
        # Each individual step is computed with small noise
        step_result = step["delta"] + np.random.randn() * 0.5
        running_total += step_result
    return round(running_total)


# Create word problems of varying complexity
def make_problem(n_steps):
    deltas = np.random.randint(-20, 20, size=n_steps)
    steps = [{"delta": int(d)} for d in deltas]
    true_answer = int(deltas.sum())
    return steps, true_answer


np.random.seed(0)

step_counts = [2, 3, 4, 5, 6, 8, 10]
n_trials = 200
results = {}

for n_steps in step_counts:
    direct_correct = 0
    cot_correct = 0

    for trial in range(n_trials):
        problem_steps, true_answer = make_problem(n_steps)
        tol = max(1, abs(true_answer) // 10 + 1)

        direct_answer = solve_direct(problem_steps)
        cot_answer = solve_chain_of_thought(problem_steps)

        if abs(direct_answer - true_answer) <= tol:
            direct_correct += 1
        if abs(cot_answer - true_answer) <= tol:
            cot_correct += 1

    results[n_steps] = {
        "direct": direct_correct / n_trials,
        "cot": cot_correct / n_trials,
    }

import numpy as np


def emergent_performance(
    log_params, threshold, sharpness, baseline=0.1, ceiling=0.9
):
    """
    Models emergent capability: near-random performance below threshold,
    rapid improvement around threshold, saturation above it.
    """
    return baseline + (ceiling - baseline) / (
        1.0 + np.exp(-sharpness * (log_params - threshold))
    )


def smooth_performance(log_params, slope=0.08, baseline=0.3):
    """
    Models a gradually scaling capability with no sharp threshold.
    """
    return np.clip(baseline + slope * log_params, 0, 1)


log_params = np.linspace(
    7, 12, 300
)  # from ~10M to ~160B parameters (log10 scale)

# Three tasks: emergent at different thresholds, and one smooth scaler
tasks = {
    "Task A (emergent, threshold=9)": emergent_performance(
        log_params, threshold=9.0, sharpness=3.0
    ),
    "Task B (emergent, threshold=10)": emergent_performance(
        log_params, threshold=10.0, sharpness=4.0
    ),
    "Task C (smooth scaling)": smooth_performance(log_params),
}

import re


# Simulated tools available to the agent
def calculator(expression):
    """Safely evaluate a simple arithmetic expression using operator decomposition."""
    try:
        # Support only simple two-operand arithmetic to avoid arbitrary execution
        match = re.fullmatch(
            r"\s*([+-]?\d+(?:\.\d+)?(?:e[+-]?\d+)?)\s*"
            r"([+\-*/])\s*"
            r"([+-]?\d+(?:\.\d+)?(?:e[+-]?\d+)?)\s*",
            expression,
        )
        if match:
            a, op, b = (
                float(match.group(1)),
                match.group(2),
                float(match.group(3)),
            )
            if op == "+":
                return a + b
            if op == "-":
                return a - b
            if op == "*":
                return a * b
            if op == "/" and b != 0:
                return a / b
    except Exception:
        pass
    return None


def lookup_fact(query):
    """Simulated knowledge base lookup."""
    facts = {
        "population of earth": 8_100_000_000,
        "speed of light km/s": 299_792,
        "pi": 3.14159265,
        "avogadro number": 6.022e23,
    }
    return facts.get(query.lower().strip())


def agent_solve(problem, available_tools):
    """
    Simple tool-use agent.
    Parses structured tool calls from the problem string and executes them.
    """
    tool_calls = []
    for tool_name, pattern, func in available_tools:
        for match_obj in re.finditer(pattern, problem, re.IGNORECASE):
            arg = match_obj.group(1).strip()
            result = func(arg)
            if result is not None:
                tool_calls.append(
                    {
                        "tool": tool_name,
                        "arg": arg,
                        "result": result,
                    }
                )
    return tool_calls


# Define tool dispatch
tools = [
    ("calculator", r"CALC\(([^)]+)\)", calculator),
    ("lookup", r"LOOKUP\(([^)]+)\)", lookup_fact),
]

# Sample problems with embedded tool calls
problems = [
    "Speed of light: LOOKUP(speed of light km/s) km/s. "
    "Seconds per year: CALC(3600 * 24) times CALC(365 * 1.0) = distance per year.",
    "Earth population: LOOKUP(population of earth) people. "
    "Divided equally among CALC(195 * 1.0) countries.",
    "Pi value from lookup: LOOKUP(pi). Area of circle r=5: CALC(78.54 * 1.0).",
]