feat(m2): multimodal splice [属性][时间戳][TS tok][事件][问题][回答] (T2.3)
- src/tsmm/model/multimodal.py: MultimodalSplicer — binds LLM input embedding layer, assembles text embeds (via tokenizer) + TS token embeds (from Projector) into inputs_embeds/attention_mask/labels; labels = -100 everywhere except answer+EOS (training); left-truncation to max_len=1024. - tests/test_multimodal.py: 8 tests against real Qwen2.5-0.5B tokenizer + stub embedding (shapes, all-ones mask, answer-only labels, TS token count in sequence, no-NaN, truncation). 74 tests passing. - configs/llm.yaml: offline model path (downloaded via proxy to ~/models/Qwen2.5-0.5B-Instruct; HF hub unreachable on host). Spec clarifications (small tier, comet-build Step 4): - Training appends answer + EOS to the prompt layout. - Build embedding layer with len(tokenizer) (incl. added special tokens), not tokenizer.vocab_size which excludes them (eos=151645).
This commit is contained in:
@@ -0,0 +1,142 @@
|
||||
"""Multimodal splice (T2.3): assemble TS tokens + text tokens into the LLM input.
|
||||
|
||||
Prompt layout (per design §2.1):
|
||||
|
||||
[属性][时间戳][TS tokens][事件][问题] # inference / prompt
|
||||
[属性][时间戳][TS tokens][事件][问题][回答][EOS] # training
|
||||
|
||||
This module produces, for a single sample, the LLM-ready tensors:
|
||||
- ``inputs_embeds`` [1, seq, hidden] (text via the bound embedding layer,
|
||||
TS tokens via the Projector output)
|
||||
- ``attention_mask`` [1, seq]
|
||||
- ``labels`` [1, seq] (-100 everywhere except the answer+EOS span)
|
||||
|
||||
Text segments are tokenized with ``add_special_tokens=False`` so we control the
|
||||
structure; the Qwen2.5 tokenizer is used as-is. Batching / variable-length
|
||||
padding is the collator's job (T2.5).
|
||||
|
||||
Truncation: if the assembled sequence exceeds ``max_len``, it is left-truncated
|
||||
to keep the TS tokens + answer (the part the model must read/predict), mirroring
|
||||
how long prompts are usually handled. The TS token block is never truncated.
|
||||
|
||||
Design ref: ``docs/superpowers/specs/2026-06-29-ts-as-modality-design.md`` §2.1.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Sequence
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpliceOutput:
|
||||
inputs_embeds: Tensor # [1, seq, hidden]
|
||||
attention_mask: Tensor # [1, seq] (0/1 int)
|
||||
labels: Tensor # [1, seq] (-100 or token id)
|
||||
|
||||
|
||||
class MultimodalSplicer:
|
||||
"""Assemble a multimodal sample from text segments + projected TS tokens.
|
||||
|
||||
The text embedding layer (typically ``llm.get_input_embeddings()``) must be
|
||||
bound via :meth:`bind` before calling :meth:`splice`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer,
|
||||
hidden_size: int,
|
||||
max_len: int = 1024,
|
||||
) -> None:
|
||||
self.tokenizer = tokenizer
|
||||
self.hidden_size = hidden_size
|
||||
self.max_len = max_len
|
||||
self._embed: Optional[torch.nn.Module] = None
|
||||
|
||||
# -- binding ----------------------------------------------------------
|
||||
def bind(self, embed_layer: torch.nn.Module) -> None:
|
||||
"""Bind the LLM text-embedding layer (input embeddings)."""
|
||||
self._embed = embed_layer
|
||||
|
||||
@property
|
||||
def embed(self) -> torch.nn.Module:
|
||||
if self._embed is None:
|
||||
raise RuntimeError("Splicer has no bound embedding layer; call bind().")
|
||||
return self._embed
|
||||
|
||||
# -- internals --------------------------------------------------------
|
||||
def _tok(self, text: str) -> Tensor:
|
||||
"""Tokenize a text segment (no special tokens) → 1D LongTensor."""
|
||||
if text is None:
|
||||
text = ""
|
||||
ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
|
||||
return torch.tensor(ids, dtype=torch.long)
|
||||
|
||||
def _embed_ids(self, ids: Tensor) -> Tensor:
|
||||
return self.embed(ids) # [len, hidden]
|
||||
|
||||
# -- public API -------------------------------------------------------
|
||||
def splice(
|
||||
self,
|
||||
*,
|
||||
attributes: str,
|
||||
timestamps: str,
|
||||
ts_embeds: Tensor, # [n_patches, hidden] (may be [0, hidden])
|
||||
events: str,
|
||||
question: str,
|
||||
answer: Optional[str] = None,
|
||||
) -> SpliceOutput:
|
||||
tok_eos = self.tokenizer.eos_token_id
|
||||
|
||||
# 1. tokenize text segments
|
||||
attr_ids = self._tok(attributes)
|
||||
ts_ids = self._tok(timestamps)
|
||||
event_ids = self._tok(events)
|
||||
q_ids = self._tok(question)
|
||||
ans_ids = self._tok(answer) if answer is not None else None
|
||||
|
||||
# 2. embed the text segments
|
||||
attr_emb = self._embed_ids(attr_ids)
|
||||
ts_text_emb = self._embed_ids(ts_ids)
|
||||
if ts_embeds.shape[0] > 0:
|
||||
ts_tok_emb = ts_embeds
|
||||
else:
|
||||
ts_tok_emb = torch.zeros(0, self.hidden_size)
|
||||
event_emb = self._embed_ids(event_ids)
|
||||
q_emb = self._embed_ids(q_ids)
|
||||
|
||||
# answer (training only)
|
||||
if ans_ids is not None:
|
||||
ans_with_eos = torch.cat([ans_ids, torch.tensor([tok_eos], dtype=torch.long)])
|
||||
ans_emb = self._embed_ids(ans_with_eos)
|
||||
|
||||
# 3. assemble embeddings along the sequence dimension
|
||||
prompt_parts = [attr_emb, ts_text_emb, ts_tok_emb, event_emb, q_emb]
|
||||
if ans_ids is not None:
|
||||
prompt_parts.append(ans_emb)
|
||||
embeds = torch.cat(prompt_parts, dim=0) # [seq, hidden]
|
||||
|
||||
# 4. build labels (default -100)
|
||||
seq_len = embeds.shape[0]
|
||||
labels = torch.full((seq_len,), -100, dtype=torch.long)
|
||||
|
||||
if ans_ids is not None:
|
||||
ans_start = seq_len - ans_emb.shape[0]
|
||||
labels[ans_start:] = ans_with_eos
|
||||
|
||||
# 5. truncate (left) to max_len, never dropping the answer
|
||||
if seq_len > self.max_len:
|
||||
keep = self.max_len
|
||||
embeds = embeds[-keep:]
|
||||
labels = labels[-keep:]
|
||||
|
||||
# 6. attention mask (single unpadded sample → all ones)
|
||||
attn = torch.ones(embeds.shape[0], dtype=torch.long)
|
||||
|
||||
return SpliceOutput(
|
||||
inputs_embeds=embeds.unsqueeze(0), # [1, seq, hidden]
|
||||
attention_mask=attn.unsqueeze(0), # [1, seq]
|
||||
labels=labels.unsqueeze(0), # [1, seq]
|
||||
)
|
||||
Reference in New Issue
Block a user