From ad7f714d62e2d0acad0799eed39d7823aba4b422 Mon Sep 17 00:00:00 2001 From: Chukwuma Nwaugha <20521315+nwaughachukwuma@users.noreply.github.com> Date: Sat, 29 Nov 2025 13:57:00 +0000 Subject: [PATCH] hfrunner.classify should return list[list[float]] not list[str] (#29671) Signed-off-by: Chukwuma Nwaugha --- tests/conftest.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 11c573befb2d2..317b36ba6cb80 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -459,14 +459,17 @@ class HfRunner: embeddings.append(embedding) return embeddings - def classify(self, prompts: list[str]) -> list[str]: + def classify(self, prompts: list[str]) -> list[list[float]]: # output is final logits all_inputs = self.get_inputs(prompts) - outputs = [] + outputs: list[list[float]] = [] problem_type = getattr(self.config, "problem_type", "") for inputs in all_inputs: output = self.model(**self.wrap_device(inputs)) + + assert isinstance(output.logits, torch.Tensor) + if problem_type == "regression": logits = output.logits[0].tolist() elif problem_type == "multi_label_classification":