| """Prompt parsing and CLIP token alignment for structured diffusion guidance.""" |
|
|
| from __future__ import annotations |
|
|
| import re |
| from dataclasses import dataclass |
| from functools import lru_cache |
| from typing import Protocol, cast |
|
|
| import torch |
| from jaxtyping import Float |
| from nltk.tree import Tree |
|
|
|
|
| class _StanzaWord(Protocol): |
| start_char: int | None |
| end_char: int | None |
|
|
|
|
| class _StanzaSentence(Protocol): |
| constituency: str |
| words: list[_StanzaWord] |
|
|
|
|
| class _StanzaDocument(Protocol): |
| sentences: list[_StanzaSentence] |
|
|
|
|
| class _StanzaPipeline(Protocol): |
| def __call__(self, text: str) -> _StanzaDocument: ... |
|
|
|
|
| class _Tokenizer(Protocol): |
| def tokenize(self, text: str) -> list[str]: ... |
|
|
|
|
| @dataclass(frozen=True) |
| class NounPhrase: |
| """A parsed noun phrase with character offsets into normalized prompt text.""" |
|
|
| text: str |
| start: int |
| end: int |
|
|
|
|
| def normalize_prompt(prompt: str) -> str: |
| """Normalize casing and trailing punctuation before parsing and token alignment.""" |
| return prompt.lower().strip().strip(".").strip() |
|
|
|
|
| def constituency_noun_phrases(prompt: str) -> list[NounPhrase]: |
| """Parse a prompt and return its multi-token noun phrases. |
| |
| The full prompt is represented separately by the pipeline. Returned offsets |
| refer to normalized prompt text so they can be aligned to CLIP BPE tokens. |
| Stanza models must be installed before this function runs. |
| """ |
| document = _constituency_pipeline()(prompt) |
| if not document.sentences: |
| return [] |
|
|
| sentence = document.sentences[0] |
| tree = Tree.fromstring(str(sentence.constituency)) |
| word_offsets: list[tuple[int, int]] = [] |
| for word in sentence.words: |
| if word.start_char is None or word.end_char is None: |
| raise ValueError("The constituency parser did not return word offsets.") |
|
|
| word_offsets.append((word.start_char, word.end_char)) |
|
|
| if len(tree.leaves()) != len(word_offsets): |
| raise ValueError( |
| "The constituency tree and parser word offsets have different lengths." |
| ) |
|
|
| phrases: list[NounPhrase] = [] |
|
|
| def visit(node: Tree | str, word_start: int) -> int: |
| if isinstance(node, str): |
| return 1 |
|
|
| cursor = word_start |
| for child in node: |
| cursor += visit(child, cursor) |
|
|
| if node.label() == "NP" and cursor - word_start > 1: |
| char_start = word_offsets[word_start][0] |
| char_end = word_offsets[cursor - 1][1] |
| text = prompt[char_start:char_end].strip() |
| if text and text != prompt: |
| phrases.append(NounPhrase(text=text, start=char_start, end=char_end)) |
|
|
| return cursor - word_start |
|
|
| visit(tree, 0) |
| return phrases |
|
|
|
|
| def scene_graph_noun_phrases(prompt: str) -> list[NounPhrase]: |
| """Return entity and relation spans from the optional scene-graph parser.""" |
| try: |
| import sng_parser |
| except ImportError as exc: |
| raise ImportError( |
| "Install structure-diffusion[scene-graph] to use parser_type='scene_graph'." |
| ) from exc |
|
|
| word_offsets = [ |
| (match.start(), match.end()) for match in re.finditer(r"\S+", prompt) |
| ] |
| graph = sng_parser.parse(prompt) |
| spans: list[tuple[int, int]] = [] |
|
|
| for entity in graph["entities"]: |
| start, end = entity["span_bounds"] |
| if end - start > 1 and 0 <= start < end <= len(word_offsets): |
| spans.append((start, end)) |
|
|
| for relation in graph["relations"]: |
| subject_start, subject_end = graph["entities"][relation["subject"]][ |
| "span_bounds" |
| ] |
| object_start, object_end = graph["entities"][relation["object"]]["span_bounds"] |
| start = min(subject_start, object_start) |
| end = max(subject_end, object_end) |
| if end - start > 1 and 0 <= start < end <= len(word_offsets): |
| spans.append((start, end)) |
|
|
| phrases: list[NounPhrase] = [] |
| seen: set[tuple[int, int]] = set() |
| for start, end in spans: |
| if (start, end) in seen: |
| continue |
|
|
| seen.add((start, end)) |
| char_start = word_offsets[start][0] |
| char_end = word_offsets[end - 1][1] |
| text = prompt[char_start:char_end].strip() |
| if text and text != prompt: |
| phrases.append(NounPhrase(text=text, start=char_start, end=char_end)) |
|
|
| return phrases |
|
|
|
|
| def token_span_from_char_span( |
| tokenizer: _Tokenizer, |
| prompt: str, |
| start: int, |
| end: int, |
| sequence_length: int, |
| ) -> tuple[int, int] | None: |
| """Map prompt character offsets to a BOS-shifted, truncated CLIP token span.""" |
| if start < 0 or end <= start or end > len(prompt): |
| raise ValueError("Phrase character offsets must be a non-empty prompt span.") |
|
|
| |
| start_token = len(tokenizer.tokenize(prompt[:start])) + 1 |
| end_token = len(tokenizer.tokenize(prompt[:end])) + 1 |
| end_token = min(end_token, sequence_length - 1) |
| if start_token >= end_token: |
| return None |
|
|
| return start_token, end_token |
|
|
|
|
| def align_phrase_context( |
| prompt_context: Float[torch.Tensor, "batch tokens dim"], |
| phrase_context: Float[torch.Tensor, "batch tokens dim"], |
| token_span: tuple[int, int], |
| phrase_token_count: int, |
| ) -> Float[torch.Tensor, "batch tokens dim"]: |
| """Copy only actual phrase-token embeddings into the matching prompt span.""" |
| start, end = token_span |
| span_length = end - start |
| if start < 1 or end > prompt_context.shape[1] or span_length < 1: |
| raise ValueError("The aligned token span is outside the prompt context.") |
|
|
| if phrase_token_count < span_length: |
| raise ValueError("The phrase encoding has fewer real tokens than its span.") |
|
|
| if phrase_context.shape[0] != prompt_context.shape[0]: |
| raise ValueError("Prompt and phrase context batches must match.") |
|
|
| aligned = prompt_context.clone() |
| |
| aligned[:, start:end] = phrase_context[:, 1 : 1 + span_length] |
| return aligned |
|
|
|
|
| @lru_cache(maxsize=1) |
| def _constituency_pipeline() -> _StanzaPipeline: |
| """Build a reusable English constituency parser without downloading models.""" |
| import stanza |
|
|
| return cast( |
| _StanzaPipeline, |
| stanza.Pipeline( |
| lang="en", |
| processors="tokenize,pos,constituency", |
| download_method=None, |
| use_gpu=False, |
| verbose=False, |
| ), |
| ) |
|
|
