docs/retrieval_contract.md

Retrieval Contract - Stage 2 (Candidate Generation / Grounding)

Stage 2 builds a high-recall candidate pool from a fixed canonical tag vocabulary. It does not invent tags and does not make final selection decisions.

Primary implementation: psq_rag/retrieval/psq_retrieval.py::psq_candidates_from_rewrite_phrases.

---

Inputs

• rewrite_phrases: Sequence[str]
• allow_nsfw_tags: bool
• context_tags: Optional[Sequence[str]] = None
• context_tag_weight: float = 1.0
• context_weight: float = 0.5
• per_phrase_k: int = 50
• per_phrase_final_k: int = 1
• global_k: int = 300
• min_tag_count: int = 0
• verbose: bool = False
• return_phrase_ranks: bool = False

Notes:

• In app runtime, these are usually overridden by env-backed settings:
  • PSQ_RETRIEVAL_GLOBAL_K (default 300)
  • PSQ_RETRIEVAL_PER_PHRASE_K (default 10)
  • PSQ_RETRIEVAL_PER_PHRASE_FINAL_K (default 1)
  • PSQ_MIN_TAG_COUNT (default 100 in app path)
• Stage 2 may be called with structural/probe tags in context_tags to improve context scoring.

---

Phrase Normalization

Each incoming phrase is normalized to:

• lowercase
• underscores converted to spaces
• trimmed and whitespace-collapsed

Then phrases are deduplicated in first-seen order.

---

Head-Term Expansion

For each normalized multi-token phrase, Stage 2 may add the last token ("head") as an extra phrase.

A head is added only when:

• phrase has at least 2 tokens
• head length is at least 3
• head is not in the built-in stopword list (and, or, the, a, an, of, to, in, on, ...).

Final phrase list = deduped original phrases + valid heads (deduped again, order-preserving).

---

Canonical Projection Rules

All candidate generation is projected into canonical tags via _project_to_canonicals(token):

1. If token itself appears as canonical (token in tag_counts or token has TF-IDF row), keep token.
2. Else if token is an alias (token in alias2tags), map to all alias targets.
3. Else drop token.

min_tag_count filtering is applied during this projection.

---

Per-Phrase Candidate Generation

For each phrase:

1. Build lookup = phrase.replace(" ", "_").
2. Compute projected_lookup = _project_to_canonicals(lookup).
3. FastText neighbors:
   • normally: fasttext_model.most_similar(lookup, topn=per_phrase_k)
   • special case:
     • if PSQ_SKIP_FASTTEXT_FOR_EXACT_ALIAS is enabled (default "1") and projected_lookup is non-empty, neighbor call is skipped for that phrase.
4. For each neighbor token:
   • project to canonical tags
   • apply NSFW filter if needed
   • keep best FastText similarity per canonical tag
   • keep alias_token that achieved that best similarity.
5. Exact-match injection:
   • each canonical tag in projected_lookup is injected with score_fasttext = 1.0
   • alias_token = lookup.

This creates a per-phrase map: tag -> score_fasttext (+ best token for debug).

---

Context Vector and Context Scores

Stage 2 builds one request-level pseudo TF-IDF vector from:

• normalized final phrases (lookup form, underscore)
• optional context_tags weighted by context_tag_weight.

Vector path:

• TF-IDF sparse vector -> SVD transform -> L2 normalize.

If zero norm:

• query_has_context = False
• all candidates use FastText-only combined score.

If non-zero:

• compute cosine vs normalized reduced TF-IDF row vectors for tags that have TF-IDF rows.
• tags with no TF-IDF row initially get score_context = None.

---

Missing-Context Imputation

When query_has_context = True, per phrase:

• gather non-None context scores for candidates in that phrase
• default_context_for_phrase = 10th percentile (q=0.10) of those scores
  • if none exist, default is 0.0.
• tags missing context receive this default and are marked context_imputed=True (verbose report only).

---

Score Fusion

Per phrase candidate:

• if no context vector:
  • score_combined = score_fasttext
• else:
  • score_combined = (1 - context_weight) * score_fasttext + context_weight * score_context

---

Per-Phrase Truncation and Required Tags

Per phrase:

• sort by score_combined desc
• keep top per_phrase_final_k.

Required tags:

• tags from projected_lookup for that phrase.
• each required tag must be present in that phrase's final list.
• if insertion is needed and list is full, evict lowest-ranked non-required item (or last item if all are required).
• re-sort by combined score after enforcement.

---

Merge Across Phrases

All per-phrase survivors are merged into a global map by canonical tag:

• sources: union of phrases that contributed this tag
• score_fasttext: max observed across phrases
• score_context: max observed across phrases (None-aware)
• score_combined: max observed across phrases
• count: from corpus tag counts

Then:

• sort by score_combined desc
• cut to global_k
• return ordered List[Candidate].

---

Output Shapes

Candidate fields:

• tag: str
• score_combined: float
• score_fasttext: Optional[float]
• score_context: Optional[float]
• count: Optional[int]
• sources: List[str]

Return variants:

• default: List[Candidate]
• verbose: (List[Candidate], phrase_reports)
• return_phrase_ranks only: (List[Candidate], phrase_rank_by_tag)
• verbose + return_phrase_ranks: (List[Candidate], phrase_reports, phrase_rank_by_tag)

phrase_reports rows include:

• phrase, normalized, lookup, tfidf_vocab, oov_terms
• candidate rows with tag, alias_token, score_fasttext, score_context, score_combined, context_imputed, count.

---

Policy Filtering

NSFW filtering uses get_nsfw_tags() from psq_rag.retrieval.state. Current source is word_rating_probabilities.csv with threshold 0.95.

---

Stage Boundary

Stage 2 is retrieval grounding only. Final tag choice happens in Stage 3 closed-set selection.