Class PromptSegment

examples/server/server.py:3850–3912

@dataclass
class PromptSegment:
    @dataclass
    class Media:
        embeddings: np.ndarray
        positions: np.ndarray
        non_causal: bool = False

    kind: Literal["text", "image", "audio", "video"]
    start_pos: int
    n_pos: int
    identity_tokens: List[int]
    decode_start_pos: int
    decode_n_pos: int
    text_tokens: List[int] = field(default_factory=list)
    media: Optional["PromptSegment.Media"] = None

    @property
    def end_pos(self) -> int:
        return self.start_pos + self.n_pos

    @property
    def batch_rows(self) -> int:
        if self.kind != "text":
            if self.media is None:
                return 0
            return int(self.media.embeddings.shape[0])
        return len(self.text_tokens)

    @property
    def decoder_position_increments(self) -> List[int]:
        if not self.identity_tokens:
            return []
        if self.kind == "text":
            return [1] * len(self.identity_tokens)
        return [*([0] * (len(self.identity_tokens) - 1)), self.decode_n_pos]

    def rows_for_capacity(self, offset: int, capacity: int) -> int:
        if capacity <= 0:
            return 0
        return min(capacity, self.end_pos - self.start_pos - offset)

    def media_slice(
        self,
        row_offset: int,
        row_count: int,
    ) -> Tuple[np.ndarray, np.ndarray, List[int]]:
        if self.media is None:
            raise RuntimeError("media segment is missing embeddings or positions")
        row_start = row_offset
        row_end = row_offset + row_count
        embeddings = self.media.embeddings[row_start:row_end]
        if len(self.media.positions) == self.batch_rows:
            positions = self.media.positions[row_start:row_end]
        else:
            positions = (
                self.media.positions.reshape(4, self.batch_rows)[:, row_start:row_end]
                .reshape(-1)
            )