Puts embedding score code in score_utils to avoid duplicated code

Signed-off-by: Gabriel Marinho <[email protected]>

Author: gmarinho2

vllm/entrypoints/llm.py

@@ -7,7 +7,6 @@
                     Tuple, Type, Union, cast, overload)
 
 import cloudpickle
-import torch
 import torch.nn as nn
 from tqdm import tqdm
 from typing_extensions import TypeVar, deprecated
@@ -25,6 +24,7 @@
                                          apply_mistral_chat_template,
                                          parse_chat_messages,
                                          resolve_chat_template_content_format)
+from vllm.entrypoints.score_utils import _cosine_similarity
 from vllm.inputs import PromptType, SingletonPrompt, TextPrompt, TokensPrompt
 from vllm.inputs.parse import is_token_prompt, parse_and_batch_prompt
 from vllm.logger import init_logger
@@ -1010,40 +1010,25 @@ def _embedding_score(
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
     ) -> List[ScoringRequestOutput]:
 
-        encoded_output = self.encode(
+        encoded_output: List[PoolingRequestOutput] = self.encode(
             text_1 + text_2,
             use_tqdm=use_tqdm,
             lora_request=lora_request,
             prompt_adapter_request=prompt_adapter_request)
-        encoded_output_1 = encoded_output[0:len(text_1)]
-        encoded_output_2 = encoded_output[len(text_1):]
+
+        encoded_output_1: List[PoolingRequestOutput] = encoded_output[
+            0:len(text_1)]
+        encoded_output_2: List[PoolingRequestOutput] = encoded_output[
+            len(text_1):]
 
         if len(encoded_output_1) == 1:
             encoded_output_1 = encoded_output_1 * len(encoded_output_2)
 
-        output_pairs = [(t1, t2)
-                        for t1, t2 in zip(encoded_output_1, encoded_output_2)]
-
-        scores = []
-        scorer = torch.nn.CosineSimilarity(0)
+        scores: List[PoolingRequestOutput] = []
 
-        for embed_1, embed_2 in output_pairs:
-            pair_score = scorer(embed_1.outputs.data, embed_2.outputs.data)
-
-            if (pad_token_id := getattr(tokenizer, "pad_token_id",
-                                        None)) is not None:
-                tokens = embed_1.prompt_token_ids + [
-                    pad_token_id
-                ] + embed_2.prompt_token_ids
-            else:
-                tokens = embed_1.prompt_token_ids + embed_2.prompt_token_ids
-
-            scores.append(
-                PoolingRequestOutput(
-                    request_id=f"{embed_1.request_id}_{embed_2.request_id}",
-                    outputs=pair_score,
-                    prompt_token_ids=tokens,
-                    finished=True))
+        scores = _cosine_similarity(tokenizer=tokenizer,
+                                    embed_1=encoded_output_1,
+                                    embed_2=encoded_output_2)
 
         items = self.engine_class.validate_outputs(scores,
                                                    PoolingRequestOutput)

vllm/entrypoints/openai/serving_score.py

@@ -3,7 +3,6 @@
 import time
 from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Union
 
-import torch
 from fastapi import Request
 
 from vllm.config import ModelConfig
@@ -16,6 +15,7 @@
                                               ScoreResponseData, UsageInfo)
 from vllm.entrypoints.openai.serving_engine import OpenAIServing
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
+from vllm.entrypoints.score_utils import _cosine_similarity
 from vllm.inputs.data import TokensPrompt
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
@@ -121,26 +121,18 @@ async def _embedding_score(
         if len(emb_text_1) == 1:
             emb_text_1 = emb_text_1 * len(emb_text_2)
 
-        scorer = torch.nn.CosineSimilarity(0)
+        embeddings_1: List[PoolingRequestOutput] = []
+        embeddings_2: List[PoolingRequestOutput] = []
 
         for emb_1, emb_2 in zip(emb_text_1, emb_text_2):
             assert emb_1 is not None
             assert emb_2 is not None
-            pair_score = scorer(emb_1.outputs.data, emb_2.outputs.data)
+            embeddings_1.append(emb_1)
+            embeddings_2.append(emb_2)
 
-            padding = []
-            if (pad_token_id := getattr(tokenizer, "pad_token_id",
-                                        None)) is not None:
-                padding = [pad_token_id]
-
-            tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
-
-            final_res_batch.append(
-                PoolingRequestOutput(
-                    request_id=f"{emb_1.request_id}_{emb_2.request_id}",
-                    outputs=pair_score,
-                    prompt_token_ids=tokens,
-                    finished=True))
+        final_res_batch = _cosine_similarity(tokenizer=tokenizer,
+                                             embed_1=embeddings_1,
+                                             embed_2=embeddings_2)
 
         return final_res_batch
 

vllm/entrypoints/score_utils.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import List, Union
+
+from torch.nn import CosineSimilarity
+
+from vllm.outputs import PoolingRequestOutput
+from vllm.transformers_utils.tokenizer import (PreTrainedTokenizer,
+                                               PreTrainedTokenizerFast)
+
+
+def _cosine_similarity(
+    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+    embed_1: List[PoolingRequestOutput],
+    embed_2: List[PoolingRequestOutput],
+) -> List[PoolingRequestOutput]:
+
+    scorer = CosineSimilarity(0)
+    scores: Union[List[PoolingRequestOutput]] = []
+
+    for emb_1, emb_2 in zip(embed_1, embed_2):
+        pair_score = scorer(emb_1.outputs.data, emb_2.outputs.data)
+
+        padding = []
+        if (pad_token_id := getattr(tokenizer, "pad_token_id",
+                                    None)) is not None:
+            padding = [pad_token_id]
+
+        tokens = emb_1.prompt_token_ids + padding + emb_2.prompt_token_ids
+
+        scores.append(
+            PoolingRequestOutput(
+                request_id=f"{emb_1.request_id}_{emb_2.request_id}",
+                outputs=pair_score,
+                prompt_token_ids=tokens,
+                finished=True))
+
+    return scores