Revert "[V1] Scatter and gather placeholders in the model runner" (#16075)
This commit is contained in:
@@ -108,46 +108,16 @@ class PromptUpdateDetails(Generic[_S]):
|
||||
full: _S
|
||||
"""The full content."""
|
||||
|
||||
is_embed: Optional[Callable[["_BoundPromptSequence"], torch.Tensor]] = None
|
||||
features: _S
|
||||
"""
|
||||
Given :attr:`full`, return a boolean mask of shape `(len(full),)`
|
||||
indicating which positions of `full` to assign embeddings to.
|
||||
|
||||
`None` (default) means to assign embeddings to all positions of `full`.
|
||||
|
||||
The embeddings are obtained by calling
|
||||
:class:`SupportsMultiModal.get_multimodal_embeddings`.
|
||||
The part of the content that corresponds to feature placeholders;
|
||||
this will be replaced by the output of the vision encoder during model
|
||||
inference.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def from_seq(seq: _S) -> "PromptUpdateDetails[_S]":
|
||||
return PromptUpdateDetails(full=seq)
|
||||
|
||||
@staticmethod
|
||||
def select_text(
|
||||
seq: _S,
|
||||
embed_text: str,
|
||||
) -> "PromptUpdateDetails[_S]":
|
||||
|
||||
def is_embed(full: "_BoundPromptSequence") -> torch.Tensor:
|
||||
embed_token_ids = encode_tokens(full.tokenizer, embed_text)
|
||||
|
||||
return torch.isin(
|
||||
torch.tensor(full.token_ids),
|
||||
torch.tensor(embed_token_ids),
|
||||
)
|
||||
|
||||
return PromptUpdateDetails(full=seq, is_embed=is_embed)
|
||||
|
||||
@staticmethod
|
||||
def select_token_id(
|
||||
seq: _S,
|
||||
embed_token_id: int,
|
||||
) -> "PromptUpdateDetails[_S]":
|
||||
return PromptUpdateDetails(
|
||||
full=seq,
|
||||
is_embed=lambda f: torch.tensor(f.token_ids) == embed_token_id,
|
||||
)
|
||||
return PromptUpdateDetails(full=seq, features=seq)
|
||||
|
||||
|
||||
PromptUpdateInfo = Union[PromptSeq, PromptUpdateDetails]
|
||||
@@ -436,7 +406,7 @@ class _BoundPromptSequence:
|
||||
@dataclass
|
||||
class _BoundPromptContent:
|
||||
full: _BoundPromptSequence
|
||||
is_embed: Optional[Callable[["_BoundPromptSequence"], torch.Tensor]]
|
||||
features: _BoundPromptSequence
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -496,8 +466,10 @@ class BoundPromptUpdate:
|
||||
|
||||
bound_full = _BoundPromptSequence.from_seq(self.tokenizer,
|
||||
content.full)
|
||||
bound_features = _BoundPromptSequence.from_seq(self.tokenizer,
|
||||
content.features)
|
||||
bound_content = _BoundPromptContent(full=bound_full,
|
||||
is_embed=content.is_embed)
|
||||
features=bound_features)
|
||||
|
||||
if cache_key is not None:
|
||||
self._content_cache[cache_key] = bound_content
|
||||
@@ -633,19 +605,15 @@ class PlaceholderFeaturesInfo:
|
||||
item_idx: int
|
||||
start_idx: int
|
||||
tokens: list[int]
|
||||
is_embed: Optional[torch.Tensor]
|
||||
|
||||
@property
|
||||
def length(self) -> int:
|
||||
return len(self.tokens)
|
||||
|
||||
def to_range(self) -> PlaceholderRange:
|
||||
# TODO: Is it worth it to optimize this by stripping the
|
||||
# leading and ending positions where `is_embed=False`?
|
||||
return PlaceholderRange(
|
||||
offset=self.start_idx,
|
||||
length=self.length,
|
||||
is_embed=self.is_embed,
|
||||
)
|
||||
|
||||
|
||||
@@ -838,17 +806,22 @@ def _iter_placeholders(
|
||||
continue
|
||||
|
||||
if prompt[start_idx:end_idx_full] == content_tokens_full:
|
||||
content_is_embed = content.is_embed
|
||||
if content_is_embed is not None:
|
||||
content_is_embed = content_is_embed(content.full)
|
||||
content_tokens_feat = content.features.token_ids
|
||||
|
||||
yield PlaceholderFeaturesInfo(
|
||||
modality=modality,
|
||||
item_idx=item_idx,
|
||||
start_idx=start_idx,
|
||||
tokens=content_tokens_full,
|
||||
is_embed=content_is_embed,
|
||||
)
|
||||
try:
|
||||
match = next(
|
||||
iter_token_matches(content_tokens_full,
|
||||
content_tokens_feat))
|
||||
yield PlaceholderFeaturesInfo(
|
||||
modality=modality,
|
||||
item_idx=item_idx,
|
||||
start_idx=start_idx + match.start_idx,
|
||||
tokens=content_tokens_feat,
|
||||
)
|
||||
except StopIteration:
|
||||
raise AssertionError(
|
||||
f"{content_tokens_feat=} should be a "
|
||||
f"subsequence of {content_tokens_full=}") from None
|
||||
|
||||
# Exclude overlapping matches
|
||||
start_idx = end_idx_full
|
||||
|
||||
Reference in New Issue
Block a user