docs/retrieval_contract.md

# Retrieval Contract -- Stage 2 (Retrieval Grounding / Candidate Generation)

Stage 2 performs **retrieval grounding** over a **closed vocabulary** of canonical e621-style tags.
It does not "tag images", and it does not do free-form generation. Its job is to produce a high-recall,
inspectable candidate pool for downstream **closed-set selection**.

---

## Inputs

- `rewrite_phrases: list[str]`
  - Output of Stage 1 query rewriting (comma-separated "tag-shaped" phrases).
  - Not canonical tags. Not underscored. High recall is preferred.

- `allow_nsfw_tags: bool`
  - If false, filter out tags in the project's `nsfw_tags` set.

- `verbose: bool`
  - If true, return per-phrase debug reports.

---

## Normalization and phrase expansion

1) Normalize rewrite phrases for internal processing:
- lowercase
- strip leading/trailing whitespace
- collapse internal whitespace to a single space

2) Treat the phrase list as a **set** (dedupe after normalization).

3) **Head-noun expansion**:
- For each multi-token phrase, add its head noun (last token) as an additional phrase.
- Apply the same set semantics so duplicates are processed once.

Example:
- Input phrases: `["big shirt", "grey shirt"]`
- Final phrase set: `{"big shirt", "grey shirt", "shirt"}`

---

## Candidate generation per phrase (FastText neighbors + canonicalization)

For each phrase `p` in the final phrase set:

1) Convert to lookup form:
- `lookup = p.replace(" ", "_")`

2) Retrieve neighbors using FastText:
- `neighbors = fasttext.most_similar(lookup, topn=per_phrase_k)`
- Note: FastText neighbors may include alias tokens and other non-canonical strings.

3) **Project neighbor tokens to canonical tags** (alias -> canonical):
- If a neighbor token is already a canonical tag (token is in `tag_counts` OR token has a TF-IDF row in `tag_to_row_index`), it maps to itself.
- Else if it is an alias, map it via `alias2tags[token]` (may map to multiple canonical tags).
- Else, drop it (not in closed vocabulary).

4) **Deduplicate by canonical tag** within this phrase:
- Keep the canonical tag with the highest FastText similarity among all tokens that mapped to it.
- Record the token that achieved that max similarity as `alias_token` for verbose reporting ("best token wins").

5) **Exact-match injection**:
- Project the phrase's own `lookup` through the same projection logic.
- For each canonical tag produced by that projection, inject it into the candidate set with:
  - `score_fasttext = 1.0`
  - `alias_token = lookup`
- This ensures the phrase canonical appears even though `most_similar()` often does not return the query token itself.

6) Apply NSFW filtering (if `allow_nsfw_tags=False`):
- Drop candidate canonical tags that are present in `nsfw_tags`.

Result: for each phrase, we have a set of canonical candidate tags with:
- `score_fasttext`
- `alias_token` (token that produced the best FastText score for that canonical tag)

---

## Context similarity (TF-IDF -> SVD cosine)

Stage 2 computes one **query context vector** for the entire request:

1) Build a pseudo TF-IDF vector from the **final phrase set** (deduped + head nouns):
- Convert each phrase to underscore form (same `lookup` rule).
- Terms that exist in the TF-IDF vocabulary (underscore lookups) contribute `(term_count * idf(term))`.
- OOV terms contribute nothing (but may be reported in verbose mode).

2) Project to SVD space and L2-normalize:
- `query_vec = normalize(svd.transform(tfidf_vec))`

If the query vector has zero norm (no recognized TF-IDF terms), then `query_has_context = False` and:
- `score_context = None` for all candidates
- `score_combined = score_fasttext` (FastText-only)

If `query_has_context = True`, compute per-candidate cosine similarity when possible:
- For tags that have a TF-IDF/SVD row: `score_context_by_tag[tag] = dot(query_vec, reduced_matrix_norm[row])`
- For tags that lack a TF-IDF/SVD row: initial `score_context = None` (may be imputed per-phrase)

### Missing context policy (per-phrase, q=0.10)

If `query_has_context = True` and a candidate tag has `score_context = None`:
- For that phrase, compute `default_context_for_phrase` as the 10th percentile (q=0.10) of the available (non-None) context scores among that phrase's candidates.
- If there are no available context scores for that phrase, `default_context_for_phrase = 0.0`.
- Impute missing context scores using `default_context_for_phrase` and mark:
  - `context_imputed = True`
Otherwise:
- `context_imputed = False`

---

## Score fusion (FastText + Context)

Compute a fused score per phrase candidate:

- If `query_has_context = False`:
  - `score_combined = score_fasttext`

- Else:
  - `score_combined = (1 - context_weight) * score_fasttext + context_weight * score_context`
  - (`score_context` may be imputed as described above)

---

## Per-phrase truncation and must-include rule

After scoring candidates for a phrase:
- Sort by `score_combined` descending.
- Keep top `per_phrase_final_k` (typically 10).

**Must-include rule (pinned exact phrase tags)**:
- Let `required_tags` be the canonical tag(s) produced by projecting the phrase's own `lookup` (`projected_lookup`).
- Each required tag must appear in that phrase's final top `per_phrase_final_k` list, even if its fused score would otherwise place it below the cutoff.
- If the list is full, evict the lowest-ranked tag that is *not* required.
- Note: `required_tags` may contain multiple canonicals if `alias2tags` maps a token to multiple tags.

This rule applies **only to the phrase's own required tags**. It does not inject tags into other phrases' lists.

---

## Merge across phrases (global candidate pool)

A canonical tag may appear in multiple per-phrase top-K lists. Stage 2 deduplicates tags into a single global record.

- `sources` is the union of phrases whose per-phrase lists contained the tag.
- `score_fasttext` is the maximum FastText score observed for the tag across those phrases.
- `score_context` is the maximum context cosine observed for the tag across those phrases (with `None` treated as missing).
- `score_combined` is the maximum fused score observed for the tag across those phrases.

Note:
- These maxima may come from different phrases; the global candidate row does not necessarily correspond to any single phrase's row.
- For tags with a TF-IDF row, `score_context` is phrase-invariant. Differences across phrases only arise for tags whose context score was imputed.

Finally:
- Sort global candidates by `score_combined` descending.
- Return top `global_k` candidates (and optionally all candidates if the app needs them).

---

## Output schema

### Stage 2 return (non-verbose)
- `candidates: list[Candidate]` (ordered)
  - `tag: str` (canonical)
  - `score_combined: float`
  - `score_fasttext: float | None`
  - `score_context: float | None` (None only when `query_has_context=False` or when missing)
  - `count: int | None`
  - `sources: list[str]`

### Optional per-phrase debug report (verbose)
For each phrase:
- `phrase: str`
- `normalized: str`
- `lookup: str`
- `tfidf_vocab: bool` (lookup is in TF-IDF vocabulary)
- `oov_terms: list[str]`
- `candidates: list[CandidateRow]` (top per-phrase list)
  - `tag: str`
  - `alias_token: str`
  - `score_fasttext: float`
  - `score_context: float | None`
  - `score_combined: float`
  - `context_imputed: bool`
  - `count: int | None`

---

## Determinism and performance constraints

- Artifact loading is **lazy** (load-on-first-use, cached thereafter).
- No feature flags for old/new behavior: delete old code paths.
- Logging must be read-only and must not affect results.

---

## NSFW tag source

- `nsfw_tags` is sourced from `word_rating_probabilities.csv` with `NSFW_THRESHOLD=0.95` as implemented in `psq_rag.retrieval.state`.