xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-19 07:27:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

simplify the returned value from generate_beam_search

Browse Source 
> includes cleaning up wrap_device, generate_prompt_perplexity and get_image_asset

Signed-off-by: Chukwuma Nwaugha <nwaughac@gmail.com>
This commit is contained in:
Chukwuma Nwaugha 2025-12-11 18:31:27 +00:00
parent 6af70e11a0
commit 95af481818
1 changed files with 11 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
tests/conftest.py
View File

@ -58,8 +58,7 @@ from vllm.distributed import (
initialize_model_parallel, initialize_model_parallel,
) )
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.logprobs import Logprob from vllm.logprobs import LogprobsOnePosition
from vllm.multimodal.base import MediaWithBytes
from vllm.multimodal.utils import fetch_image from vllm.multimodal.utils import fetch_image
from vllm.outputs import RequestOutput from vllm.outputs import RequestOutput
from vllm.sampling_params import BeamSearchParams from vllm.sampling_params import BeamSearchParams
@ -93,8 +92,7 @@ PromptVideoInput = _PromptMultiModalInput[np.ndarray]
def _read_prompts(filename: str) -> list[str]: def _read_prompts(filename: str) -> list[str]:
with open(filename) as f: with open(filename) as f:
prompts = f.readlines() return f.readlines()
return prompts
class ImageAssetPrompts(TypedDict): class ImageAssetPrompts(TypedDict):
@ -267,7 +265,7 @@ class HfRunner:
if isinstance(x, dict): if isinstance(x, dict):
return {k: self.wrap_device(v, device) for k, v in x.items()} return {k: self.wrap_device(v, device) for k, v in x.items()}
if hasattr(x, "device") and x.device.type == device: if hasattr(x.device, "type") and x.device.type == device:
return x return x
return x.to(device) return x.to(device)
@ -993,8 +991,8 @@ class VllmRunner:
perplexities = [] perplexities = []
for output in outputs: for output in outputs:
output = cast(TokensTextLogprobsPromptLogprobs, output) assert isinstance(output, TokensTextLogprobsPromptLogprobs)
token_datas = cast(list[dict[int, Logprob] | None], output[3]) token_datas = cast(list[LogprobsOnePosition | None], output[3])
assert token_datas[0] is None assert token_datas[0] is None
token_log_probs = [] token_log_probs = []
for token_data in token_datas[1:]: for token_data in token_datas[1:]:
@ -1025,12 +1023,11 @@ class VllmRunner:
BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens), BeamSearchParams(beam_width=beam_width, max_tokens=max_tokens),
concurrency_limit=concurrency_limit, concurrency_limit=concurrency_limit,
) )
returned_outputs = []
for output in outputs: return [
token_ids = [x.tokens for x in output.sequences] ([x.tokens for x in output.sequences], [x.text for x in output.sequences])
texts = [x.text for x in output.sequences] for output in outputs
returned_outputs.append((token_ids, texts)) ]
return returned_outputs
def classify(self, prompts: list[str]) -> list[list[float]]: def classify(self, prompts: list[str]) -> list[list[float]]:
req_outputs = self.llm.classify(prompts) req_outputs = self.llm.classify(prompts)
@ -1405,11 +1402,7 @@ class LocalAssetServer:
return f"{self.base_url}/{name}" return f"{self.base_url}/{name}"
def get_image_asset(self, name: str) -> Image.Image: def get_image_asset(self, name: str) -> Image.Image:
image = fetch_image(self.url_for(name)) return fetch_image(self.url_for(name))
# Unwrap MediaWithBytes if present
if isinstance(image, MediaWithBytes):
image = image.media
return image
@pytest.fixture(scope="session") @pytest.fixture(scope="session")