File size: 6,539 Bytes
f675a6c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | """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.")
# CLIP reserves token index 0 for BOS, so shift text-token offsets by one.
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()
# Skip the phrase encoding's BOS token and copy only the aligned phrase tokens.
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,
),
)
|