| 12105 | self.batch.n_tokens += 1 |
| 12106 | |
| 12107 | def add_batch_embeddings( |
| 12108 | self, |
| 12109 | *, |
| 12110 | seq_id: int, |
| 12111 | embeddings: np.ndarray, |
| 12112 | positions: np.ndarray, |
| 12113 | output_indices: Sequence[Optional[int]], |
| 12114 | ) -> None: |
| 12115 | if self.batch.n_tokens: |
| 12116 | raise RuntimeError("cannot mix token and embedding batches") |
| 12117 | if self._embedding_batch is not None: |
| 12118 | raise RuntimeError("only one embedding batch is supported per scheduler step") |
| 12119 | embeddings = np.ascontiguousarray(embeddings, dtype=np.float32) |
| 12120 | positions = np.ascontiguousarray(positions, dtype=np.int32).reshape(-1) |
| 12121 | n_tokens = int(embeddings.shape[0]) |
| 12122 | if n_tokens == 0: |
| 12123 | return |
| 12124 | if embeddings.ndim != 2 or embeddings.shape[1] != self.n_embd_inp: |
| 12125 | raise RuntimeError("embedding batch shape does not match model input embedding size") |
| 12126 | if len(positions) not in {n_tokens, n_tokens * 4}: |
| 12127 | raise RuntimeError("embedding position length mismatch") |
| 12128 | if len(output_indices) != n_tokens: |
| 12129 | raise RuntimeError("embedding output index length mismatch") |
| 12130 | pos_array = (llama_cpp.llama_pos * len(positions))( |
| 12131 | *[int(pos) for pos in positions] |
| 12132 | ) |
| 12133 | n_seq_id_array = (ctypes.c_int32 * n_tokens)(*[1] * n_tokens) |
| 12134 | seq_id_array = (llama_cpp.llama_seq_id * 1)(llama_cpp.llama_seq_id(seq_id)) |
| 12135 | seq_ids_array = (ctypes.POINTER(llama_cpp.llama_seq_id) * (n_tokens + 1))() |
| 12136 | for index in range(n_tokens): |
| 12137 | seq_ids_array[index] = seq_id_array |
| 12138 | logits_array = (ctypes.c_int8 * n_tokens)( |
| 12139 | *[int(output_index is not None) for output_index in output_indices] |
| 12140 | ) |
| 12141 | batch = llama_cpp.llama_batch( |
| 12142 | n_tokens=n_tokens, |
| 12143 | token=None, |
| 12144 | embd=embeddings.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), |
| 12145 | pos=pos_array, |
| 12146 | n_seq_id=n_seq_id_array, |
| 12147 | seq_id=seq_ids_array, |
| 12148 | logits=logits_array, |
| 12149 | ) |
| 12150 | self._embedding_batch = batch |
| 12151 | self._embedding_batch_refs = [ |
| 12152 | embeddings, |
| 12153 | positions, |
| 12154 | pos_array, |
| 12155 | n_seq_id_array, |
| 12156 | seq_id_array, |
| 12157 | seq_ids_array, |
| 12158 | logits_array, |
| 12159 | ] |
| 12160 | |
| 12161 | def decode(self) -> None: |
| 12162 | batch = self._embedding_batch if self._embedding_batch is not None else self.batch |