A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 08:30 UTC
hfrunner.classify should return list[list[float]] not list[str] (#29671)
Signed-off-by: Chukwuma Nwaugha <nwaughac@gmail.com>
Chukwuma Nwaugha committed
Nov 29, 2025 at 13:57 UTC
ad7f714d62e2d0acad0799eed39d7823aba4b422
1 file changed
+5
-2
tests/conftest.py
+5
-2
@@ -459,14 +459,17 @@ class HfRunner:
459
embeddings.append(embedding)
460
return embeddings
461
462
- def classify(self, prompts: list[str]) -> list[str]:
462
+ def classify(self, prompts: list[str]) -> list[list[float]]:
463
# output is final logits
464
all_inputs = self.get_inputs(prompts)
465
- outputs = []
465
+ outputs: list[list[float]] = []
466
problem_type = getattr(self.config, "problem_type", "")
467
468
for inputs in all_inputs:
469
output = self.model(**self.wrap_device(inputs))
470
+
471
+ assert isinstance(output.logits, torch.Tensor)
472
+
473
if problem_type == "regression":
474
logits = output.logits[0].tolist()
475
elif problem_type == "multi_label_classification":