Repository: vllm-project/vllm
File: tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py (Python)

Code before:
            self._build_render_params(prompt),
        )
        return engine_prompts

    def post_process(
        self,
        model_output: Sequence[PoolingRequestOutput],
        request_id: str | None = None,
        **kwargs,
    ) -> IOProcessorOutput:
        num_prompt_tokens = 0
        response_data = []
        for idx in range(len(model_output)):
            mo = model_output[idx]
            sparse_embedding = {}
            num_prompt_tokens += len(mo.prompt_token_ids)
            if len(mo.prompt_token_ids) != len(mo.outputs.data):
                # this is the case that add_special_tokens is True,
                # which means first token and last token are special tokens
                mo.prompt_token_ids = mo.prompt_token_ids[1:]
            for token_id, weight in zip(mo.prompt_token_ids, mo.outputs.data):
                if token_id not in sparse_embedding:
                    sparse_embedding[token_id] = weight
                    continue
                if weight > sparse_embedding[token_id]:
                    sparse_embedding[token_id] = weight
            response_data.append(
                SparseEmbeddingResponseData(
                    index=idx, sparse_embedding=sparse_embedding
                )
            )
        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            total_tokens=num_prompt_tokens,
        )
        resp = SparseEmbeddingResponse(
            request_id=request_id,
            data=response_data,
            usage=usage,
        )

        return resp

    def output_to_response(
        self, plugin_output: IOProcessorOutput
    ) -> IOProcessorResponse:
        return IOProcessorResponse(
            request_id=plugin_output.request_id,
            data=plugin_output,
        )


Review comment (suggestion):
```suggestion
                sparse_embedding[token_id] = max(weight, sparse_embedding.get(token_id, 0.0))
```

Code after:

    def post_process(
        self,
        model_output: Sequence[PoolingRequestOutput],
        request_id: str | None = None,
        **kwargs,
    ) -> SparseEmbeddingResponse:
        num_prompt_tokens = 0
        response_data = []
        return_tokens = self._get_sparse_embedding_request(request_id).return_tokens
        for idx in range(len(model_output)):
            mo = model_output[idx]
            sparse_embedding: dict[int, float] = {}
            num_prompt_tokens += len(mo.prompt_token_ids)
            if len(mo.prompt_token_ids) != len(mo.outputs.data):
                # this is the case that add_special_tokens is True,
                # which means first token and last token are special tokens
                mo.prompt_token_ids = mo.prompt_token_ids[1:]
            for token_id, weight in zip(mo.prompt_token_ids, mo.outputs.data.tolist()):
                sparse_embedding[token_id] = max(
                    weight, sparse_embedding.get(token_id, 0.0)
                )
            response_data.append(
                SparseEmbeddingResponseData(
                    index=idx,
                    sparse_embedding=self._build_sparse_embedding_token_weights(
                        sparse_embedding,
                        return_tokens,
                    ),
                )
            )

        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            total_tokens=num_prompt_tokens,
        )
        resp = SparseEmbeddingResponse(
            data=response_data,
            usage=usage,
        )

        return resp