torch.optim
Optimizers.
SGD
Stochastic gradient descent.
"""
def test_chunk_key_stable_and_distinct():
assert chunk_key(UNIT) == chunk_key(dict(UNIT))
other = {**UNIT, "anchor": "other"}
assert chunk_key(other) != chunk_key(UNIT)
def test_batches_splits_exactly():
items = list(range(10))
got = list(batches(items, 4))
assert got == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]]
def test_query_prefix_is_bge_convention():
assert QUERY_PREFIX.startswith("Represent this sentence")
def test_symbol_from_url():
from index.embed import symbol_from_url
url = "https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.sdpa.html"
assert symbol_from_url(url) == "torch.nn.functional.sdpa"
tut = "https://docs.pytorch.org/tutorials/beginner/intro.html"
assert symbol_from_url(tut) == ""
def test_indexed_text_prepends_symbol_and_heading():
from index.embed import indexed_text
unit = {
"url": "https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html",
"heading_path": ["torch.optim.SGD", "Parameters"],
"content": "lr (float) — learning rate.",
}
text = indexed_text(unit)
assert text.startswith("torch.optim.SGD")
assert "torch.optim.SGD > Parameters" in text
assert "lr (float)" in text
def test_indexed_text_prepends_the_extracted_synopsis():
# the page's own docstring summary (extracted at chunk time, no LLM) leads
# every chunk right after the symbol — the free half of the semantic bridge
from index.embed import indexed_text
unit = {
"url": "https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html",
"heading_path": ["CrossEntropyLoss", "Parameters"],
"synopsis": "This criterion computes the cross entropy loss between "
"input logits and target.",
"content": "weight (Tensor) — a manual rescaling weight.",
}
text = indexed_text(unit)
assert text.index("torch.nn.CrossEntropyLoss") < text.index("This criterion")
assert text.index("This criterion") < text.index("CrossEntropyLoss > Parameters")
def test_indexed_text_prepends_the_pages_gloss_when_present(monkeypatch):
# Contextual Retrieval: the plain-language gloss is the semantic bridge
# between a signature-shaped reference page and descriptive questions —
# it must land in indexed_text (→ both the vector and the tsvector)
from index.embed import indexed_text
url = "https://docs.pytorch.org/docs/stable/generated/torch.optim.SGD.html"
gloss = "Stochastic gradient descent optimizer with optional momentum."
monkeypatch.setattr("index.embed.load_glosses", lambda: {url: gloss})
unit = {"url": url, "heading_path": ["torch.optim.SGD"], "content": "lr (float)"}
text = indexed_text(unit)
assert gloss in text
assert text.index("torch.optim.SGD") < text.index(gloss) < text.index("lr (float)")
# a page with no gloss is unchanged
assert gloss not in indexed_text({**unit, "url": url.replace("SGD", "Adam")})
def test_gloss_stamp_changes_the_recipe_with_gloss_content(monkeypatch, tmp_path):
# any change to the committed gloss file must force the one-time full
# re-embed by itself — a stale recipe would silently keep old vectors
from index import embed
monkeypatch.setattr(embed, "GLOSSES_PATH", tmp_path / "glosses.jsonl")
assert embed._gloss_stamp() == "none"
embed.GLOSSES_PATH.write_text('{"url": "u", "gloss": "g"}\n')
first = embed._gloss_stamp()
embed.GLOSSES_PATH.write_text('{"url": "u", "gloss": "different"}\n')
assert first not in ("none", embed._gloss_stamp())
def test_indexed_text_folds_in_the_pages_hypothetical_questions(monkeypatch):
# QuOTE-style: the question phrasing itself joins both channels, so a
# descriptive query matches question→question instead of question→signature
from index.embed import indexed_text
url = "https://docs.pytorch.org/docs/stable/generated/torch.nn.Linear.html"
questions = [
"What's the standard fully-connected layer that applies a weight matrix and bias?",
"How do I add a dense layer to my network?",
]
monkeypatch.setattr("index.embed.load_glosses", lambda: {})
monkeypatch.setattr("index.embed.load_questions", lambda: {url: questions})
unit = {"url": url, "heading_path": ["Linear"], "content": "in_features (int)"}
text = indexed_text(unit)
for q in questions:
assert q in text
# questions sit between the symbol block and the body, like the gloss
assert text.index("torch.nn.Linear") < text.index(questions[0]) < text.index("in_features")
# a page with no questions is unchanged
assert questions[0] not in indexed_text({**unit, "url": url.replace("Linear", "GELU")})
def test_questions_stamp_changes_the_recipe_with_question_content(monkeypatch, tmp_path):
# same contract as the gloss stamp: new/changed questions force the re-embed
from index import embed
monkeypatch.setattr(embed, "QUESTIONS_PATH", tmp_path / "questions.jsonl")
assert embed._questions_stamp() == "none"
embed.QUESTIONS_PATH.write_text('{"url": "u", "questions": ["q"]}\n')
first = embed._questions_stamp()
embed.QUESTIONS_PATH.write_text('{"url": "u", "questions": ["other"]}\n')
assert first not in ("none", embed._questions_stamp())
def test_iter_corpus_units_walks_snapshot(tmp_path):
save_page("https://docs.pytorch.org/docs/stable/optim.html", "core", HTML, tmp_path)
units = list(iter_corpus_units(tmp_path))
assert units, "expected chunks from the snapshot page"
assert all(u["url"].endswith("optim.html") for u in units)
assert any("SGD" in " > ".join(u["heading_path"]) for u in units)
def test_chunk_key_part_zero_keeps_the_legacy_format():
# backward compatibility is the whole point: introducing parts must not
# change existing rows' keys, or the next build re-embeds the entire corpus
from index.embed import chunk_key
unit = {"url": "https://x/p.html", "anchor": "a", "heading_path": ["A", "B"]}
assert chunk_key({**unit, "part": 0}) == chunk_key(unit)
assert chunk_key({**unit, "part": 1}) != chunk_key(unit)
assert chunk_key({**unit, "part": 1}) != chunk_key({**unit, "part": 2})