Method decode_batch

llama_cpp/llama.py:1108–1134

(seq_sizes: List[int])

        data: Union[List[List[float]], List[List[List[float]]]] = []

        def decode_batch(seq_sizes: List[int]):
            self._ctx.kv_cache_clear()
            self._ctx.decode(self._batch)
            self._batch.reset()

            # store embeddings
            if pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE:
                pos: int = 0
                for i, size in enumerate(seq_sizes):
                    ptr = llama_cpp.llama_get_embeddings(self._ctx.ctx)
                    embedding: List[List[float]] = [
                        ptr[pos + j * n_embd : pos + (j + 1) * n_embd]
                        for j in range(size)
                    ]
                    if normalize:
                        embedding = [
                            internals.normalize_embedding(e) for e in embedding
                        ]
                    data.append(embedding)
                    pos += size
            else:
                for i in range(len(seq_sizes)):
                    ptr = llama_cpp.llama_get_embeddings_seq(self._ctx.ctx, i)
                    embedding: List[float] = ptr[:n_embd]
                    if normalize:
                        embedding = internals.normalize_embedding(embedding)
                    data.append(embedding)

        # init state
        total_tokens = 0