M02: Tokens & Context Limits

# pip install tiktoken
import tiktoken

# Load the cl100k_base encoding (GPT-4 / close to Mistral)
enc = tiktoken.get_encoding("cl100k_base")

# ── Basic token counting ────────────────────────────────────────
samples = [
    "Hello, world!",
    "tokenization",
    "don't use pseudocode",
    "ChatGPT is one token",
    "from openai import OpenAI",
]

for text in samples:
    ids = enc.encode(text)
    tokens = [enc.decode_single_token_bytes(t).decode("utf-8", errors="replace")
              for t in ids]
    print(f"{len(ids):3d} tokens | {tokens!r}")
    print(f"         text: {text!r}")
    print()

# ── Chat message overhead ──────────────────────────────────────
# OpenAI-style chat adds 3 tokens per message + 3 for the reply primer
def count_chat_tokens(messages: list[dict]) -> int:
    """Count tokens for a list of chat messages including overhead."""
    total = 3  # primer for assistant reply
    for msg in messages:
        total += 4  # role, content field markers, sep
        total += len(enc.encode(msg.get("role", "")))
        total += len(enc.encode(msg.get("content", "")))
    return total

messages = [
    {"role": "system", "content": "You are a helpful coding assistant."},
    {"role": "user",   "content": "Explain what a token is in 2 sentences."},
]
print(f"Chat token estimate: {count_chat_tokens(messages)}")
# Output: Chat token estimate: ~31

# pip install transformers
from transformers import AutoTokenizer

# Downloads ~500KB of config from HuggingFace Hub (cached)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")

samples = [
    "Hello, world!",
    "tokenization",
    "don't use pseudocode",
    "ChatGPT is one token",
    "from openai import OpenAI",
]

import tiktoken
enc = tiktoken.get_encoding("cl100k_base")

print(f"{'Text':<35} {'tiktoken':>8} {'Mistral':>8} {'diff':>6}")
print("-" * 62)
for text in samples:
    tik_n = len(enc.encode(text))
    mis_n = len(tokenizer.encode(text, add_special_tokens=False))
    diff = mis_n - tik_n
    sign = "+" if diff > 0 else ""
    print(f"{text!r:<35} {tik_n:>8} {mis_n:>8} {sign+str(diff):>6}")

# pip install tiktoken openai
from __future__ import annotations
import tiktoken
from openai import OpenAI

# ─────────────────────────────────────────────────────────────────
# CHUNK 1: PromptBudget class definition
# ─────────────────────────────────────────────────────────────────
class PromptBudget:
    """Track and manage token usage across a conversation."""

    # Per-message overhead in tiktoken cl100k_base format
    TOKENS_PER_MESSAGE = 4    # role + content markers + sep
    TOKENS_REPLY_PRIMER = 3   # assistant reply primer tokens

    def __init__(self, model: str = "mistral", max_context: int = 32_000):
        self.model = model
        self.max_context = max_context
        # Use cl100k_base as a close approximation for Mistral
        self._enc = tiktoken.get_encoding("cl100k_base")

    def estimate_tokens(self, text: str) -> int:
        """Count tokens in a plain string."""
        return len(self._enc.encode(text))

    def _count_messages(self, messages: list[dict]) -> int:
        """Count tokens for a messages list including overhead."""
        total = self.TOKENS_REPLY_PRIMER
        for msg in messages:
            total += self.TOKENS_PER_MESSAGE
            total += len(self._enc.encode(msg.get("role", "")))
            total += len(self._enc.encode(msg.get("content", "")))
        return total

    def _count_tools(self, tools: list[dict] | None) -> int:
        """Rough estimate: serialize tools to JSON and count tokens."""
        if not tools:
            return 0
        import json
        return len(self._enc.encode(json.dumps(tools)))

    def remaining(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        reserve_output: int = 512,
    ) -> int:
        """Return how many tokens are left for new messages + output."""
        used = self._count_messages(messages) + self._count_tools(tools)
        return self.max_context - used - reserve_output

    def fits(
        self,
        text: str,
        messages: list[dict],
        tools: list[dict] | None = None,
        reserve_output: int = 512,
    ) -> bool:
        """True if adding `text` as a new message still fits in budget."""
        new_msg_tokens = self.estimate_tokens(text) + self.TOKENS_PER_MESSAGE
        return self.remaining(messages, tools, reserve_output) >= new_msg_tokens

    def truncate_history(
        self,
        messages: list[dict],
        reserve_output: int = 512,
    ) -> list[dict]:
        """
        Drop oldest non-system messages until the conversation fits.
        Always preserves the first message if it's the system prompt.
        Returns a new list — does not mutate the input.
        """
        result = list(messages)
        # Identify system prompt boundary (first message, role=system)
        start = 1 if result and result[0].get("role") == "system" else 0

        while True:
            tokens_used = self._count_messages(result)
            headroom = self.max_context - tokens_used - reserve_output
            if headroom >= 0:
                break
            if len(result) <= start + 1:
                # Can't drop any more — just return what we have
                break
            result.pop(start)  # Remove oldest non-system message

        return result

# ─────────────────────────────────────────────────────────────────
# CHUNK 2: Usage example with Ollama
# ─────────────────────────────────────────────────────────────────
def main():
    budget = PromptBudget(model="mistral", max_context=32_000)
    client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

    system = {"role": "system", "content": "You are a concise coding assistant."}
    history: list[dict] = [system]

    turns = [
        "What is a Python list comprehension?",
        "Show me an example that squares even numbers from 1 to 20.",
        "How does that compare to a regular for-loop?",
        "What's the memory usage difference?",
    ]

    for user_text in turns:
        print(f"\nUser: {user_text}")
        print(f"  Tokens remaining: {budget.remaining(history)}")

        if not budget.fits(user_text, history):
            print("  [truncating history to fit]")
            history = budget.truncate_history(history)

        history.append({"role": "user", "content": user_text})

        try:
            response = client.chat.completions.create(
                model="mistral",
                messages=history,
                max_tokens=256,
            )
            reply = response.choices[0].message.content
            history.append({"role": "assistant", "content": reply})
            print(f"Assistant: {reply[:120]}...")
        except Exception as e:
            print(f"  Error: {e}")
            # Remove the user message we just added so state stays consistent
            history.pop()

if __name__ == "__main__":
    main()

# pip install openai
import time
from dataclasses import dataclass
from openai import OpenAI

# ─────────────────────────────────────────────────────────────────
# CHUNK 1: TokenTimer dataclass + wrapper
# ─────────────────────────────────────────────────────────────────
@dataclass
class TimerResult:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    elapsed_seconds: float
    decode_tokens_per_sec: float
    content: str

class TokenTimer:
    """Wrap an Ollama call and report latency + throughput metrics."""

    def __init__(self, model: str = "mistral"):
        self.model = model
        self.client = OpenAI(
            base_url="http://localhost:11434/v1",
            api_key="ollama",
        )

    def run(
        self,
        messages: list[dict],
        max_tokens: int = 256,
    ) -> TimerResult:
        """Send a chat request and return timing metrics."""
        t_start = time.perf_counter()
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                max_tokens=max_tokens,
            )
        except Exception as e:
            raise RuntimeError(f"Ollama request failed: {e}") from e

        elapsed = time.perf_counter() - t_start
        usage = response.usage
        content = response.choices[0].message.content or ""

        # decode speed = output tokens / total time
        decode_tps = (usage.completion_tokens / elapsed) if elapsed > 0 else 0

        return TimerResult(
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            total_tokens=usage.total_tokens,
            elapsed_seconds=elapsed,
            decode_tokens_per_sec=decode_tps,
            content=content,
        )

# ─────────────────────────────────────────────────────────────────
# CHUNK 2: Benchmark over three prompt sizes
# ─────────────────────────────────────────────────────────────────
def make_prompt(approx_tokens: int) -> list[dict]:
    """Generate a padding prompt of approximately the given token count."""
    word = "benchmark "
    # ~1.3 chars per token → n_chars ≈ tokens * 1.3
    padding = word * (approx_tokens // 2)
    return [
        {"role": "system", "content": "You are a helpful assistant. Answer briefly."},
        {"role": "user", "content": f"Please summarize the following:\n\n{padding}\n\nSummarize in one sentence."},
    ]

def main():
    timer = TokenTimer(model="mistral")

    # Warm-up (avoids model load time polluting benchmarks)
    print("Warming up...")
    _ = timer.run([{"role": "user", "content": "hi"}], max_tokens=5)

    print("\nBenchmark results:")
    print(f"{'Approx Input':>14} {'Prompt Tok':>10} {'Comp Tok':>9} {'Elapsed':>8} {'Decode tok/s':>13}")
    print("-" * 60)

    for target in [100, 500, 1_000]:
        msgs = make_prompt(target)
        result = timer.run(msgs, max_tokens=64)
        print(
            f"{target:>12}→ "
            f"{result.prompt_tokens:>9} "
            f"{result.completion_tokens:>9} "
            f"{result.elapsed_seconds:>7.2f}s "
            f"{result.decode_tokens_per_sec:>12.1f}"
        )

if __name__ == "__main__":
    main()