Generators and itertools

squares = (x * x for x in range(10**9))   # instant, O(1) memory
next(squares)                              # 0
next(squares)                              # 1

total = sum(len(t) for t in texts)
total = sum((len(t) for t in texts), start=0)   # extra parentheses needed with a second argument

gen = (x / norm for x in vectors)   # `vectors` evaluated now
norm = 2.0                          # `norm` looked up lazily: changes the result!

def is_safe(text: str) -> bool:
    return moderation_client.classify(text).safe   # network call

ok = all(is_safe(t) for t in outputs)       # stops at the first unsafe text
ok = all([is_safe(t) for t in outputs])     # classifies everything first

best = max((h.score for h in hits), default=0.0)
top_hit = max(hits, key=lambda h: h.score, default=None)

from typing import Iterator, Iterable

def clean_lines(lines: Iterable[str]) -> Iterator[str]:
    for line in lines:
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        yield line

def walk_messages(conversations):
    for conv in conversations:
        yield from conv["messages"]

def flatten(obj):
    """Recursively flatten nested lists/tuples (arbitrary depth)."""
    if isinstance(obj, (list, tuple)):
        for item in obj:
            yield from flatten(item)
    else:
        yield obj

list(flatten([1, [2, [3, [4]]], (5,)]))   # [1, 2, 3, 4, 5]

lines = read_lines("corpus.txt")
cleaned = clean_lines(lines)
chunks = (c for line in cleaned for c in split_into_chunks(line, 512))
for chunk in chunks:
    index.add(chunk)

def bad(path):
    with open(path) as fh:
        return (line.strip() for line in fh)   # file closes on return

def good(path):
    with open(path) as fh:
        for line in fh:
            yield line.strip()                  # file stays open while iterating

import json

def read_jsonl(path):
    with open(path, encoding="utf-8") as fh:
        for lineno, line in enumerate(fh, 1):
            if line.strip():
                try:
                    yield json.loads(line)
                except json.JSONDecodeError as e:
                    raise ValueError(f"{path}:{lineno}: {e}") from e

n_long = sum(1 for rec in read_jsonl("train.jsonl") if len(rec["text"]) > 4000)

from functools import partial

def file_chunks(path, size=1 << 20):
    with open(path, "rb") as fh:
        yield from iter(partial(fh.read, size), b"")

from itertools import islice

def batched(iterable, n):
    """Backport of itertools.batched (3.12+) for Python 3.9-3.11."""
    if n < 1:
        raise ValueError("n must be at least one")
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch

list(batched("ABCDEFG", 3))   # [('A', 'B', 'C'), ('D', 'E', 'F'), ('G',)]

for batch in batched(chunks, 100):
    vector_store.upsert(
        ids=[c.id for c in batch],
        vectors=embed([c.text for c in batch]),
        metadata=[c.metadata for c in batch],
    )

def batch_by_tokens(texts, count_tokens, max_tokens=8000, max_items=256):
    batch, used = [], 0
    for text in texts:
        n = count_tokens(text)
        if batch and (used + n > max_tokens or len(batch) == max_items):
            yield batch
            batch, used = [], 0
        batch.append(text)
        used += n
    if batch:
        yield batch

batches = [items[i:i + size] for i in range(0, len(items), size)]

from itertools import pairwise, islice
from collections import deque

deltas = [b - a for a, b in pairwise(timestamps)]          # 3.10+
bigrams = list(pairwise(tokens))

def sliding_window(iterable, n):
    """Overlapping windows of size n: ABCDE, 3 -> ABC BCD CDE."""
    it = iter(iterable)
    window = deque(islice(it, n - 1), maxlen=n)
    for x in it:
        window.append(x)
        yield tuple(window)

def overlapping_chunks(tokens, size=512, overlap=64):
    step = size - overlap
    return [tokens[i:i + size] for i in range(0, max(len(tokens) - overlap, 1), step)]

from itertools import groupby
from operator import itemgetter

rows = sorted(results, key=itemgetter("model"))
for model, group in groupby(rows, key=itemgetter("model")):
    scores = [r["score"] for r in group]       # consume before the next iteration
    print(model, sum(scores) / len(scores))

rle = [(k, sum(1 for _ in g)) for k, g in groupby("aaabccdd")]   # [('a', 3), ('b', 1), ('c', 2), ('d', 2)]

from itertools import accumulate, product, combinations
import operator

token_counts = [120, 340, 90, 410]
offsets = list(accumulate(token_counts, initial=0))      # [0, 120, 460, 550, 960]
running_max = list(accumulate(losses, min))               # best loss so far

# keep messages until the context budget is exhausted
from itertools import takewhile
kept = [m for m, total in takewhile(lambda p: p[1] <= budget,
                                    zip(messages, accumulate(len_tokens(m) for m in messages)))]

# hyper-parameter grid
grid = [
    {"lr": lr, "batch_size": bs, "warmup": w}
    for lr, bs, w in product([1e-5, 3e-5], [16, 32], [0, 500])
]

# all unordered pairs for a similarity matrix (n*(n-1)/2 pairs)
pairs = [(a.id, b.id, cosine(a.vec, b.vec)) for a, b in combinations(docs, 2)]

from itertools import chain, islice, zip_longest, tee, filterfalse
from collections import deque

def take(n, iterable):
    return list(islice(iterable, n))

def first(iterable, default=None):
    return next(iter(iterable), default)

def consume(iterator):
    deque(iterator, maxlen=0)

def partition(pred, iterable):
    """Split into (false_items, true_items) iterators."""
    t1, t2 = tee(iterable)
    return filterfalse(pred, t1), filter(pred, t2)

def unique_everseen(iterable, key=None):
    seen = set()
    for item in iterable:
        k = item if key is None else key(item)
        if k not in seen:
            seen.add(k)
            yield item

def roundrobin(*iterables):
    sentinel = object()
    for group in zip_longest(*iterables, fillvalue=sentinel):
        yield from (x for x in group if x is not sentinel)

passed, failed = [], []
for case in eval_cases:
    (passed if case.score >= 0.8 else failed).append(case)