Token Counter in Machine Learning

import tiktoken

def count_tokens_openai(text: str, model: str = "gpt-4o") -> dict:
    """Count tokens for OpenAI models using tiktoken."""
    # Get the encoding for the specified model
    encoding = tiktoken.encoding_for_model(model)
    
    # Encode the text into token IDs
    token_ids = encoding.encode(text)
    token_count = len(token_ids)
    
    # Pricing per million tokens (as of Feb 2026)
    pricing = {
        "gpt-4o": {"input": 2.50, "output": 10.00},
        "gpt-4o-mini": {"input": 0.15, "output": 0.60},
        "gpt-4.1": {"input": 2.00, "output": 8.00},
    }
    
    cost_per_million = pricing.get(model, {"input": 2.50, "output": 10.00})
    estimated_input_cost_usd = (token_count / 1_000_000) * cost_per_million["input"]
    
    return {
        "model": model,
        "encoding": encoding.name,
        "token_count": token_count,
        "token_ids": token_ids[:10],  # First 10 for inspection
        "estimated_input_cost_usd": round(estimated_input_cost_usd, 6),
        "estimated_input_cost_inr": round(estimated_input_cost_usd * 84, 4),
    }

# Example usage
text = "Explain how token counting works in large language models."
result = count_tokens_openai(text)
print(f"Token count: {result['token_count']}")
print(f"Encoding: {result['encoding']}")
print(f"Cost (USD): ${result['estimated_input_cost_usd']}")
print(f"Cost (INR): ₹{result['estimated_input_cost_inr']}")

# Compare across models
for model in ["gpt-4o", "gpt-4o-mini"]:
    r = count_tokens_openai(text, model)
    print(f"{model}: {r['token_count']} tokens, ${r['estimated_input_cost_usd']}")

from transformers import AutoTokenizer
from typing import Optional

class MultiModelTokenCounter:
    """Token counter supporting any HuggingFace-compatible model."""
    
    def __init__(self):
        self._cache: dict[str, AutoTokenizer] = {}
    
    def _get_tokenizer(self, model_name: str) -> AutoTokenizer:
        if model_name not in self._cache:
            self._cache[model_name] = AutoTokenizer.from_pretrained(model_name)
        return self._cache[model_name]
    
    def count(self, text: str, model_name: str) -> int:
        tokenizer = self._get_tokenizer(model_name)
        return len(tokenizer.encode(text))
    
    def count_with_details(self, text: str, model_name: str) -> dict:
        tokenizer = self._get_tokenizer(model_name)
        encoded = tokenizer.encode(text)
        tokens = tokenizer.convert_ids_to_tokens(encoded)
        return {
            "model": model_name,
            "token_count": len(encoded),
            "tokens_preview": tokens[:20],
            "vocab_size": tokenizer.vocab_size,
            "chars_per_token": round(len(text) / max(len(encoded), 1), 2),
        }
    
    def compare_models(self, text: str, model_names: list[str]) -> list[dict]:
        """Compare token counts across multiple models."""
        results = []
        for model in model_names:
            result = self.count_with_details(text, model)
            results.append(result)
        return sorted(results, key=lambda x: x["token_count"])

# Usage
counter = MultiModelTokenCounter()

text = "मुंबई में आज का मौसम बहुत अच्छा है।"  # Hindi text
models = [
    "meta-llama/Llama-3.1-8B",
    "mistralai/Mistral-7B-v0.1",
    "google/gemma-2-2b",
]

for result in counter.compare_models(text, models):
    print(f"{result['model']}: {result['token_count']} tokens "
          f"({result['chars_per_token']} chars/token)")

import tiktoken
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class TokenBudget:
    """Represents the token allocation for a single LLM request."""
    context_window: int  # Total model context window
    max_output_tokens: int  # Reserved for response
    system_prompt_tokens: int = 0
    user_query_tokens: int = 0
    context_tokens: int = 0
    history_tokens: int = 0
    
    @property
    def total_input_tokens(self) -> int:
        return (self.system_prompt_tokens + self.user_query_tokens + 
                self.context_tokens + self.history_tokens)
    
    @property
    def available_tokens(self) -> int:
        return self.context_window - self.max_output_tokens - self.total_input_tokens
    
    @property
    def is_within_budget(self) -> bool:
        return self.available_tokens >= 0
    
    @property
    def utilization_pct(self) -> float:
        usable = self.context_window - self.max_output_tokens
        return round((self.total_input_tokens / usable) * 100, 1) if usable > 0 else 0.0


class TokenBudgetManager:
    """Manages token budgets for LLM requests."""
    
    MODEL_CONFIGS = {
        "gpt-4o": {"window": 128_000, "encoding": "o200k_base"},
        "gpt-4o-mini": {"window": 128_000, "encoding": "o200k_base"},
        "gpt-4.1": {"window": 1_047_576, "encoding": "o200k_base"},
    }
    
    def __init__(self, model: str = "gpt-4o", max_output_tokens: int = 4096):
        config = self.MODEL_CONFIGS[model]
        self.encoding = tiktoken.get_encoding(config["encoding"])
        self.model = model
        self.budget = TokenBudget(
            context_window=config["window"],
            max_output_tokens=max_output_tokens,
        )
    
    def count(self, text: str) -> int:
        return len(self.encoding.encode(text))
    
    def allocate_system_prompt(self, text: str) -> int:
        count = self.count(text)
        self.budget.system_prompt_tokens = count
        return count
    
    def allocate_user_query(self, text: str) -> int:
        count = self.count(text)
        self.budget.user_query_tokens = count
        return count
    
    def allocate_context(self, chunks: list[str], max_tokens: Optional[int] = None) -> list[str]:
        """Add context chunks until budget or max_tokens is exhausted."""
        max_tokens = max_tokens or self.budget.available_tokens
        selected_chunks = []
        tokens_used = 0
        
        for chunk in chunks:
            chunk_tokens = self.count(chunk)
            if tokens_used + chunk_tokens > max_tokens:
                break
            selected_chunks.append(chunk)
            tokens_used += chunk_tokens
        
        self.budget.context_tokens = tokens_used
        return selected_chunks
    
    def get_budget_report(self) -> dict:
        b = self.budget
        return {
            "model": self.model,
            "context_window": b.context_window,
            "max_output_tokens": b.max_output_tokens,
            "system_prompt_tokens": b.system_prompt_tokens,
            "user_query_tokens": b.user_query_tokens,
            "context_tokens": b.context_tokens,
            "history_tokens": b.history_tokens,
            "total_input_tokens": b.total_input_tokens,
            "available_tokens": b.available_tokens,
            "utilization_pct": b.utilization_pct,
            "is_within_budget": b.is_within_budget,
        }

# Usage
manager = TokenBudgetManager(model="gpt-4o", max_output_tokens=4096)

manager.allocate_system_prompt("You are a helpful assistant for IRCTC.")
manager.allocate_user_query("What is the PNR status for ticket 2124567890?")

# Simulate RAG context chunks
chunks = [
    "PNR 2124567890: Train 12301 Rajdhani Express, Delhi to Kolkata...",
    "Current status: Confirmed, Coach B2, Berth 45, Departure 18:00...",
    "Refund policy: Cancellation before 24 hours gets 75% refund...",
    "Alternative trains: 12303 Poorva Express departs at 20:00...",
]
selected = manager.allocate_context(chunks)

report = manager.get_budget_report()
print(f"Using {report['total_input_tokens']} of {report['context_window']} tokens")
print(f"Utilization: {report['utilization_pct']}%")
print(f"Chunks included: {len(selected)} of {len(chunks)}")

import tiktoken
from typing import Generator

def chunk_by_tokens(
    text: str,
    chunk_size: int = 512,
    chunk_overlap: int = 64,
    model: str = "gpt-4o",
) -> Generator[dict, None, None]:
    """Split text into chunks of exactly `chunk_size` tokens.
    
    Unlike character-based chunking, this ensures each chunk
    is exactly the right size for the target model.
    """
    encoding = tiktoken.encoding_for_model(model)
    tokens = encoding.encode(text)
    total_tokens = len(tokens)
    
    if total_tokens <= chunk_size:
        yield {
            "text": text,
            "token_count": total_tokens,
            "chunk_index": 0,
            "start_token": 0,
            "end_token": total_tokens,
        }
        return
    
    step = chunk_size - chunk_overlap
    chunk_index = 0
    
    for start in range(0, total_tokens, step):
        end = min(start + chunk_size, total_tokens)
        chunk_tokens = tokens[start:end]
        chunk_text = encoding.decode(chunk_tokens)
        
        yield {
            "text": chunk_text,
            "token_count": len(chunk_tokens),
            "chunk_index": chunk_index,
            "start_token": start,
            "end_token": end,
        }
        
        chunk_index += 1
        if end >= total_tokens:
            break

# Usage
document = """India's Unified Payments Interface (UPI) processed over 
15 billion transactions in December 2025, handling more than 
₹23 lakh crore in value. PhonePe and Google Pay together account 
for roughly 85% of UPI transaction volume. The system's success 
has made it a model for real-time payment infrastructure globally, 
with countries like Singapore, UAE, and France adopting 
interoperable QR-code based systems inspired by UPI..."""

for chunk in chunk_by_tokens(document, chunk_size=128, chunk_overlap=16):
    print(f"Chunk {chunk['chunk_index']}: {chunk['token_count']} tokens")
    print(f"  Preview: {chunk['text'][:80]}...")
    print()

# Token counter configuration (YAML)
token_counter:
  default_model: gpt-4o
  cache_tokenizers: true
  
  models:
    gpt-4o:
      encoding: o200k_base
      context_window: 128000
      pricing:
        input_per_million: 2.50
        output_per_million: 10.00
        cached_input_per_million: 1.25
    gpt-4o-mini:
      encoding: o200k_base  
      context_window: 128000
      pricing:
        input_per_million: 0.15
        output_per_million: 0.60
    claude-sonnet-4-5:
      tokenizer: anthropic  # Use Anthropic's token counting API
      context_window: 200000
      pricing:
        input_per_million: 3.00
        output_per_million: 15.00
        cached_input_per_million: 0.30
    gemini-2.5-pro:
      tokenizer: google  # Use Google's token counting API
      context_window: 1000000
      pricing:
        input_per_million: 1.25
        output_per_million: 10.00
  
  budget:
    default_max_output_tokens: 4096
    warn_at_utilization_pct: 85
    hard_limit_utilization_pct: 95
    
  quotas:
    per_user_daily_tokens: 1000000
    per_team_monthly_tokens: 100000000
    alert_at_pct: 80