Cleanup

Author: SiDen

vllm/beam/beam.py

@@ -1,37 +1,44 @@
 from collections.abc import AsyncGenerator
+from typing import Union
+
 from vllm.beam.debug import BeamDebugInfo
 from vllm.beam.penalty import PenaltyComputer
 import torch
 from vllm.beam.ranking import RankingComputer
+from vllm.entrypoints.openai.protocol import CompletionResponse, ErrorResponse
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 class BeamScorer:
     def __init__(self, classi_idx):
         self.penalty_computer = PenaltyComputer(classi_idx)
         self.ranking_computer = RankingComputer(classi_idx)
 
-    async def collapse_beams(self, responses: list[AsyncGenerator], chunk_num = 0, max_chunks = 4):
-            debug_info = [BeamDebugInfo() for _ in responses]
-            
-            scores = torch.zeros(len(responses), dtype=torch.float)
- 
-            heads = [response.choices[0].additional_heads[0] for response in responses]
-            heads_tensor = torch.tensor(heads, dtype=torch.float)
-            if len(heads_tensor) > 0:
-                penalties = self.penalty_computer.compute(heads_tensor, debug_info)
-                scores -= penalties
-
-                ranking_scores = self.ranking_computer.compute(
+    async def pick_best_beam(self, responses: list[
+        Union[AsyncGenerator[str, None], CompletionResponse, ErrorResponse]]) -> Union[
+        AsyncGenerator[str, None], CompletionResponse, ErrorResponse]:
+        debug_info = [BeamDebugInfo() for _ in responses]
+
+        scores = torch.zeros(len(responses), dtype=torch.float)
+
+        heads = [response.choices[0].additional_heads[0] for response in responses]
+        heads_tensor = torch.tensor(heads, dtype=torch.float)
+        if len(heads_tensor) > 0:
+            penalties = self.penalty_computer.compute(heads_tensor, debug_info)
+            scores -= penalties
+
+            ranking_scores = self.ranking_computer.compute(
                 heads_tensor, debug_info
-               )
-                scores *= ranking_scores
+            )
+            scores *= ranking_scores
 
-            for i in range(len(responses)):
-                debug_info[i].final_score = scores[i]
-                debug_info[i].content = responses[i].choices[0].text
+        for i in range(len(responses)):
+            debug_info[i].final_score = scores[i]
+            debug_info[i].content = responses[i].choices[0].text
 
-            print('debug_info', debug_info)
+        logger.debug('debug_info: %s', debug_info)
 
-            best_idx = torch.argmax(scores).item()
-            return responses[best_idx]
-    
+        best_idx = torch.argmax(scores).item()
+        return responses[best_idx]

vllm/beam/filtering.py

@@ -1,10 +1,14 @@
 import asyncio
 from collections.abc import AsyncGenerator
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 from urllib.request import Request
 from vllm.beam.debug import BeamDebugInfo
 import torch
-from vllm.entrypoints.openai.protocol import CompletionRequest
+from vllm.entrypoints.openai.protocol import CompletionRequest, CompletionResponse, \
+    ErrorResponse
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 def format_filter(filter_params_list):
@@ -13,65 +17,75 @@ def format_filter(filter_params_list):
         raise ValueError(f"Expect 2 keys, seeing {filter_params_list=}")
     return [dict(zip(keys, vals)) for vals in filter_params_list]
 
+
 DEFAULT_CHAR_SERVER_FILTER = format_filter(
-        [
-            ("annotations_porn", 0.1098),
-            ("annotations_racist", 0.2814),
-            ("annotations_disturbing", 0.1827),
-            ("annotations_harmful_promotes_selfharm", 0.0749),
-            ("annotations_harmful_promotes_terrorism", 0.1129),
-        ]
-        )
-    
+    [
+        ("annotations_porn", 0.1098),
+        ("annotations_racist", 0.2814),
+        ("annotations_disturbing", 0.1827),
+        ("annotations_harmful_promotes_selfharm", 0.0749),
+        ("annotations_harmful_promotes_terrorism", 0.1129),
+    ]
+)
+
 MAX_GENERATIONS = 10
 _CHUNK_SIZE = 16
 
+
 class BeamValidator:
     def __init__(self, classi_idx, classifier_names):
         self.classi_idx = classi_idx
         self.classifier_names = classifier_names
 
-    async def get_n_valid_beams(self, create_completion: Callable, request: CompletionRequest, raw_request: Optional[Request] = None):
+    async def get_n_valid_beams(self, create_completion: Callable,
+                                request: CompletionRequest,
+                                raw_request: Optional[Request] = None) -> list[
+        Union[AsyncGenerator[str, None], CompletionResponse, ErrorResponse]]:
         request.stream = False
         n = request.n
         request.n = 1
+        # TODO(@tanuj): accept max tokens as a parameter
         request.max_tokens = _CHUNK_SIZE
         request.echo = True
         tasks = []
+        # TODO(@tanuj): deep copy request and raw_request?
         for _ in range(n):
             request = request
             tasks.append(create_completion(
                 request,
+                raw_request=raw_request,
             ))
         res = await asyncio.gather(*tasks)
         request.n = n
         beam_validator_res = self.validate(res)
         filtered_res = [r for r, valid in zip(res, beam_validator_res) if valid]
-        print('everything is filtered', len(filtered_res) == 0)
+        logger.debug("Filtered count: %d", len(filtered_res))
         if len(filtered_res) == 0:
             return res
-        
+
         return filtered_res
-    
-    def validate(self, responses: list[AsyncGenerator], debug_infos_G: list[BeamDebugInfo] = None):
-        #TODO(@tanuj) - share this with the beam scorer
+
+    def validate(self, responses: list[AsyncGenerator],
+                 debug_infos_G: list[BeamDebugInfo] = None):
+        # TODO(@tanuj) - share this with the beam scorer
         heads = [response.choices[0].additional_heads[0] for response in responses]
         heads_tensor = torch.tensor(heads, dtype=torch.float)
         prob_GC = torch.sigmoid(heads_tensor)
         valid_G = torch.ones(prob_GC.shape[0], dtype=torch.bool)
-        
+
         for g in range(heads_tensor.shape[0]):
-            filtered = self.get_filtered_classifiers(prob_GC[g], DEFAULT_CHAR_SERVER_FILTER)
+            filtered = self.get_filtered_classifiers(prob_GC[g],
+                                                     DEFAULT_CHAR_SERVER_FILTER)
             if debug_infos_G is not None:
                 debug_infos_G[g].filtered_classifiers = filtered
-            
+
             if filtered:
                 valid_G[g] = False
                 for choice in responses[g].choices:
                     choice.is_filtered = True
 
         return valid_G
-    
+
     def get_filtered_classifiers(self, prob_C, filter_params) -> list[str]:
         relevant_filters = [
             (p["name"], self.classi_idx[p["name"]], p["threshold"])
@@ -87,4 +101,4 @@ def get_filtered_classifiers(self, prob_C, filter_params) -> list[str]:
             if prob_C[idx] > threshold:
                 ret.append(name)
 
-        return ret
+        return ret