Stdlib Power Tools

from collections import ChainMap, Counter, OrderedDict, defaultdict, deque, namedtuple

# Group chunks by document
by_doc = defaultdict(list)
for chunk in chunks:
    by_doc[chunk.doc_id].append(chunk)

# Label distribution and top tokens
labels = Counter(ex["label"] for ex in dataset)
top_tokens = Counter(tok for text in corpus for tok in text.split()).most_common(20)
labels.update(["positive", "negative"])           # add counts
imbalance = labels.most_common(1)[0][1] / max(labels.total(), 1)   # .total(): 3.10+

# Bounded chat memory: keeps the last 20 turns automatically
history = deque(maxlen=20)
history.append({"role": "user", "content": "hi"})

# Simple LRU cache
class LRU:
    def __init__(self, capacity: int):
        self.capacity, self.data = capacity, OrderedDict()

    def get(self, key):
        if key not in self.data:
            return None
        self.data.move_to_end(key)
        return self.data[key]

    def put(self, key, value):
        self.data[key] = value
        self.data.move_to_end(key)
        if len(self.data) > self.capacity:
            self.data.popitem(last=False)

# Layered configuration: CLI > env > file > defaults
config = ChainMap(cli_args, env_config, file_config, DEFAULTS)
model = config["model"]

# Record type
Hit = namedtuple("Hit", ["doc_id", "score"])
best = max(hits, key=lambda h: h.score)

from typing import NamedTuple

class Hit(NamedTuple):
    doc_id: str
    score: float

from itertools import accumulate, chain, combinations, cycle, groupby, islice, product

# Peek at the first 5 records of a huge JSONL file without reading it all
with open("train.jsonl", encoding="utf-8") as f:
    sample = [json.loads(line) for line in islice(f, 5)]

# Grid search over generation params
grid = [dict(temperature=t, top_p=p) for t, p in product([0.0, 0.3, 0.7], [0.9, 1.0])]

# Pack chunks into a context window by cumulative token count
lengths = [len(enc.encode(c.text)) for c in ranked_chunks]
fits = [c for c, total in zip(ranked_chunks, accumulate(lengths)) if total <= budget]

# groupby needs input sorted by the same key
rows.sort(key=lambda r: r["doc_id"])
for doc_id, group in groupby(rows, key=lambda r: r["doc_id"]):
    merge(doc_id, list(group))

# Near-duplicate check over all pairs (O(n^2): small n only)
dupes = [(a.id, b.id) for a, b in combinations(docs, 2) if jaccard(a, b) > 0.9]

# Round-robin over API keys / replicas
endpoints = cycle(["https://a.example", "https://b.example"])
next_url = next(endpoints)

import functools
import time

@functools.cache
def get_tokenizer(name: str):
    from transformers import AutoTokenizer
    return AutoTokenizer.from_pretrained(name)

def timed(fn):
    @functools.wraps(fn)                     # keeps fn.__name__, docstring, signature for tooling
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        try:
            return fn(*args, **kwargs)
        finally:
            logger.info("%s took %.1f ms", fn.__name__, (time.perf_counter() - start) * 1000)
    return wrapper

@functools.singledispatch
def to_text(obj) -> str:
    raise TypeError(f"cannot convert {type(obj).__name__}")

@to_text.register
def _(obj: str) -> str:
    return obj

@to_text.register
def _(obj: dict) -> str:
    return json.dumps(obj, ensure_ascii=False)

embed_small = functools.partial(embed, model="text-embedding-3-small")
vectors = list(map(embed_small, texts))

from operator import attrgetter, itemgetter, methodcaller

hits.sort(key=attrgetter("score"), reverse=True)        # instead of lambda h: h.score
rows.sort(key=itemgetter("doc_id", "chunk_index"))     # multi-key
ids = list(map(itemgetter("id"), records))
stripped = list(map(methodcaller("strip"), lines))

from pathlib import Path

ROOT = Path(__file__).resolve().parent
DATA = ROOT / "data"

def iter_documents(folder: Path, exts=(".md", ".txt")):
    for path in sorted(folder.rglob("*")):
        if path.is_file() and path.suffix.lower() in exts:
            yield path.relative_to(folder).as_posix(), path.read_text(encoding="utf-8")

out = DATA / "embeddings" / f"{run_id}.npy"
out.parent.mkdir(parents=True, exist_ok=True)

import time
from contextlib import ExitStack, contextmanager, nullcontext, suppress

@contextmanager
def timer(name: str):
    start = time.perf_counter()
    try:
        yield
    finally:
        logger.info("%s: %.2fs", name, time.perf_counter() - start)

with timer("embed corpus"):
    vectors = embed_corpus(texts)

with suppress(FileNotFoundError):
    Path("index.lock").unlink()

def write_shards(paths: list[Path], rows):
    with ExitStack() as stack:
        files = [stack.enter_context(p.open("w", encoding="utf-8")) for p in paths]
        for i, row in enumerate(rows):
            files[i % len(files)].write(json.dumps(row) + "\n")

def maybe_locked(lock=None):
    return lock if lock is not None else nullcontext()

from dataclasses import asdict, dataclass, field, replace

@dataclass(frozen=True)
class RetrievalConfig:
    top_k: int = 5
    min_score: float = 0.2
    rerank: bool = True
    filters: dict = field(default_factory=dict)    # never filters: dict = {}

    def __post_init__(self):
        if self.top_k <= 0:
            raise ValueError("top_k must be positive")

base = RetrievalConfig()
strict = replace(base, min_score=0.5)            # copy with changes (immutable update)
payload = asdict(strict)                          # recursive dict for JSON/logging

from enum import Enum, Flag, auto

class FinishReason(str, Enum):          # 3.11+: StrEnum
    STOP = "stop"
    LENGTH = "length"
    TOOL_CALLS = "tool_calls"
    CONTENT_FILTER = "content_filter"

reason = FinishReason(choice.finish_reason)     # raises ValueError on unknown values
if reason is FinishReason.LENGTH:
    logger.warning("truncated output")

class Capability(Flag):
    TEXT = auto()
    VISION = auto()
    TOOLS = auto()

caps = Capability.TEXT | Capability.TOOLS
if Capability.TOOLS in caps: ...

import bisect
import heapq

# Top-k without full sort: O(n log k)
top = heapq.nlargest(10, scored_docs, key=lambda d: d.score)

# Merge already-sorted result lists from several shards
merged = list(heapq.merge(shard_a, shard_b, key=lambda d: -d.score))

# Priority queue of pending jobs
jobs: list[tuple[int, str]] = []
heapq.heappush(jobs, (priority, job_id))
priority, job_id = heapq.heappop(jobs)

# Bucket a score into bands with bisect
BANDS = [0.3, 0.6, 0.85]
LABELS = ["low", "medium", "high", "exact"]
label = LABELS[bisect.bisect_right(BANDS, score)]

import statistics
p50, p95 = statistics.median(latencies), statistics.quantiles(latencies, n=100)[94]

import hashlib
def cache_key(model: str, prompt: str) -> str:
    return hashlib.sha256(f"{model}\0{prompt}".encode()).hexdigest()