Repository: kserve/kserve

Issue:
Add option to skip postprocess in huggingface server
/kind feature

**Describe the solution you'd like**

Add option to skip the default `postprocess` method of huggingface server, as users may want to postprocess the raw prediction results their own way, which may not be possible right now.

For example, let's take the `sequence_classification` task:
https://github.com/kserve/kserve/blob/f6782431cd0751b3dc504fd51d56da75e6e12027/python/huggingfaceserver/huggingfaceserver/model.py#L352-L358

Currently, the `postprocess` makes it output only the predicted idx (i.e. `0`, or `1`). So let's say if a user wants to get the probabilities of each idx (i.e. `{0: 0.61, 1: 0.39}`), it is not possible even if they implement their own `postprocess`, because the default `postprocess` will run first.

**Anything else you would like to add:**


**Links to the design documents:**
[Optional, start with the short-form [RFC template](https://docs.google.com/document/d/1UcBeLfZ_JMGpVrPJmYtEIVH_9Y4U3AEKQdq_IKuOMrU) to outline your ideas and get early feedback.]
[Required, use the longer-form [design doc template](https://docs.google.com/document/d/1Mtoui_PP2a9N59NjYHnsvrdJ8t2iKFwIJAx1zRO_I1c) to specify and discuss your design in more detail]



Discussion:

Resolving patch:
diff --git a/python/huggingfaceserver/huggingfaceserver/__main__.py b/python/huggingfaceserver/huggingfaceserver/__main__.py
--- a/python/huggingfaceserver/huggingfaceserver/__main__.py
+++ b/python/huggingfaceserver/huggingfaceserver/__main__.py
@@ -107,6 +107,11 @@ def list_of_strings(arg):
 parser.add_argument(
     "--return_token_type_ids", action="store_true", help="Return token type ids"
 )
+parser.add_argument(
+    "--return_probabilities",
+    action="store_true",
+    help="Return all probabilities",
+)
 
 parser = maybe_add_vllm_cli_parser(parser)
 
diff --git a/python/huggingfaceserver/huggingfaceserver/encoder_model.py b/python/huggingfaceserver/huggingfaceserver/encoder_model.py
--- a/python/huggingfaceserver/huggingfaceserver/encoder_model.py
+++ b/python/huggingfaceserver/huggingfaceserver/encoder_model.py
@@ -79,6 +79,7 @@ def __init__(
         model_revision: Optional[str] = None,
         tokenizer_revision: Optional[str] = None,
         trust_remote_code: bool = False,
+        return_probabilities: bool = False,
         predictor_config: Optional[PredictorConfig] = None,
     ):
         super().__init__(model_name, predictor_config)
@@ -93,6 +94,7 @@ def __init__(
         self.model_revision = model_revision
         self.tokenizer_revision = tokenizer_revision
         self.trust_remote_code = trust_remote_code
+        self.return_probabilities = return_probabilities
 
         if model_config:
             self.model_config = model_config
@@ -259,23 +261,42 @@ def postprocess(
             num_rows, num_cols = outputs.shape
             for i in range(num_rows):
                 out = outputs[i].unsqueeze(0)
-                predicted_idx = out.argmax().item()
-                inferences.append(predicted_idx)
+                if self.return_probabilities:
+                    inferences.append(dict(enumerate(out.numpy().flatten())))
+                else:
+                    predicted_idx = out.argmax().item()
+                    inferences.append(predicted_idx)
             return get_predict_response(request, inferences, self.name)
         elif self.task == MLTask.fill_mask:
             num_rows = outputs.shape[0]
             for i in range(num_rows):
                 mask_pos = (input_ids == self._tokenizer.mask_token_id)[i]
                 mask_token_index = mask_pos.nonzero(as_tuple=True)[0]
-                predicted_token_id = outputs[i, mask_token_index].argmax(axis=-1)
-                inferences.append(self._tokenizer.decode(predicted_token_id))
+                if self.return_probabilities:
+                    probabilities = torch.softmax(outputs[i, mask_token_index], dim=-1)
+                    decoded_probabilities = []
+                    for idx, probs in enumerate(probabilities):
+                        token_probs = []
+                        for token_id, prob in enumerate(probs):
+                            token = self._tokenizer.decode([token_id])
+                            token_probs.append({f"{token}": f"{prob.item():.4f}"})
+                        decoded_probabilities.append(token_probs)
+                    inferences.append(decoded_probabilities)
+                else:
+                    predicted_token_id = outputs[i, mask_token_index].argmax(axis=-1)
+                    inferences.append(self._tokenizer.decode(predicted_token_id))
             return get_predict_response(request, inferences, self.name)
         elif self.task == MLTask.token_classification:
             num_rows = outputs.shape[0]
             for i in range(num_rows):
                 output = outputs[i].unsqueeze(0)
-                predictions = torch.argmax(output, dim=2)
-                inferences.append(predictions.tolist())
+                if self.return_probabilities:
+                    for values in output.tolist():
+                        res = [{k: v for k, v in enumerate(value)} for value in values]
+                        inferences.append([res])
+                else:
+                    predictions = torch.argmax(output, dim=2)
+                    inferences.append(predictions.tolist())
             return get_predict_response(request, inferences, self.name)
         else:
             raise ValueError(


Test patch:
diff --git a/python/huggingfaceserver/huggingfaceserver/test_model.py b/python/huggingfaceserver/huggingfaceserver/test_model.py
--- a/python/huggingfaceserver/huggingfaceserver/test_model.py
+++ b/python/huggingfaceserver/huggingfaceserver/test_model.py
@@ -22,12 +22,14 @@
 )
 from pytest_httpx import HTTPXMock
 from transformers import AutoConfig
+from pytest import approx
 
 from .task import infer_task_from_model_architecture
 from .encoder_model import HuggingfaceEncoderModel
 from .generative_model import HuggingfaceGenerativeModel
 from .task import MLTask
 import torch
+from .test_output import bert_token_classification_retrun_prob_expected_output
 
 
 @pytest.fixture(scope="module")
@@ -81,6 +83,33 @@ def bert_base_yelp_polarity():
     model.stop()
 
 
+@pytest.fixture(scope="module")
+def bert_base_return_prob():
+    model = HuggingfaceEncoderModel(
+        "bert-base-uncased-yelp-polarity",
+        model_id_or_path="textattack/bert-base-uncased-yelp-polarity",
+        task=MLTask.sequence_classification,
+        return_probabilities=True,
+    )
+    model.load()
+    yield model
+    model.stop()
+
+
+@pytest.fixture(scope="module")
+def bert_token_classification_retrun_prob():
+    model = HuggingfaceEncoderModel(
+        "bert-large-cased-finetuned-conll03-english",
+        model_id_or_path="dbmdz/bert-large-cased-finetuned-conll03-english",
+        do_lower_case=True,
+        add_special_tokens=False,
+        return_probabilities=True,
+    )
+    model.load()
+    yield model
+    model.stop()
+
+
 @pytest.fixture(scope="module")
 def bert_token_classification():
     model = HuggingfaceEncoderModel(
@@ -236,6 +265,32 @@ async def test_bert_sequence_classification(bert_base_yelp_polarity):
     assert response == {"predictions": [1, 1]}
 
 
+@pytest.mark.asyncio
+async def test_bert_sequence_classification_return_probabilities(bert_base_return_prob):
+    request = "Hello, my dog is cute."
+    response = await bert_base_return_prob(
+        {"instances": [request, request]}, headers={}
+    )
+
+    assert response == {
+        "predictions": [
+            {0: approx(-3.1508713), 1: approx(3.5892851)},
+            {0: approx(-3.1508713), 1: approx(3.589285)},
+        ]
+    }
+
+
+@pytest.mark.asyncio
+async def test_bert_token_classification_return_prob(
+    bert_token_classification_retrun_prob,
+):
+    request = "Hello, my dog is cute."
+    response = await bert_token_classification_retrun_prob(
+        {"instances": [request, request]}, headers={}
+    )
+    assert response == bert_token_classification_retrun_prob_expected_output
+
+
 @pytest.mark.asyncio
 async def test_bert_token_classification(bert_token_classification):
     request = "HuggingFace is a company based in Paris and New York"
diff --git a/python/huggingfaceserver/huggingfaceserver/test_output.py b/python/huggingfaceserver/huggingfaceserver/test_output.py
new file mode 100644
--- /dev/null
+++ b/python/huggingfaceserver/huggingfaceserver/test_output.py
@@ -0,0 +1,183 @@
+# Copyright 2024 The KServe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pytest import approx
+
+
+bert_token_classification_retrun_prob_expected_output = {
+    "predictions": [
+        [
+            [
+                {
+                    0: approx(10.125612258911133),
+                    1: approx(-1.818464756011963),
+                    2: approx(-1.3191171884536743),
+                    3: approx(-1.9324339628219604),
+                    4: approx(-1.4850239753723145),
+                    5: approx(-1.69266676902771),
+                    6: approx(-0.898107647895813),
+                    7: approx(-1.725127935409546),
+                    8: approx(0.32057392597198486),
+                },
+                {
+                    0: approx(9.647306442260742),
+                    1: approx(-2.112745523452759),
+                    2: approx(-0.8831668496131897),
+                    3: approx(-2.719135046005249),
+                    4: approx(-0.47473397850990295),
+                    5: approx(-2.2424018383026123),
+                    6: approx(0.6101505160331726),
+                    7: approx(-2.2192084789276123),
+                    8: approx(-0.654518187046051),
+                },
+                {
+                    0: approx(10.364561080932617),
+                    1: approx(-2.240158796310425),
+                    2: approx(-0.9236820340156555),
+                    3: approx(-2.623302936553955),
+                    4: approx(-0.501063346862793),
+                    5: approx(-1.9418426752090454),
+                    6: approx(-0.041013482958078384),
+                    7: approx(-2.12089204788208),
+                    8: approx(-1.2565152645111084),
+                },
+                {
+                    0: approx(9.21961784362793),
+                    1: approx(-2.1359012126922607),
+                    2: approx(-0.1689995676279068),
+                    3: approx(-3.0277929306030273),
+                    4: approx(0.2589835822582245),
+                    5: approx(-2.4426753520965576),
+                    6: approx(0.4815778136253357),
+                    7: approx(-2.3223936557769775),
+                    8: approx(-0.23837946355342865),
+                },
+                {
+                    0: approx(11.023712158203125),
+                    1: approx(-2.0757784843444824),
+                    2: approx(-0.8648976683616638),
+                    3: approx(-2.6088907718658447),
+                    4: approx(-0.781069278717041),
+                    5: approx(-1.7040152549743652),
+                    6: approx(-0.5265535712242126),
+                    7: approx(-1.7141871452331543),
+                    8: approx(-1.0127886533737183),
+                },
+                {
+                    0: approx(9.288057327270508),
+                    1: approx(-1.9795234203338623),
+                    2: approx(0.3095782697200775),
+                    3: approx(-2.781409740447998),
+                    4: approx(-0.4492073655128479),
+                    5: approx(-1.7794851064682007),
+                    6: approx(-0.3115033507347107),
+                    7: approx(-2.0755014419555664),
+                    8: approx(-0.8829556107521057),
+                },
+                {
+                    0: approx(7.968824863433838),
+                    1: approx(-2.0212578773498535),
+                    2: approx(-0.43740716576576233),
+                    3: approx(-3.6762948036193848),
+                    4: approx(-0.012759938836097717),
+                    5: approx(-2.320978879928589),
+                    6: approx(0.6871926784515381),
+                    7: approx(-2.541991949081421),
+                    8: approx(0.32606029510498047),
+                },
+            ]
+        ],
+        [
+            [
+                {
+                    0: approx(10.125612258911133),
+                    1: approx(-1.818464756011963),
+                    2: approx(-1.3191171884536743),
+                    3: approx(-1.9324339628219604),
+                    4: approx(-1.4850239753723145),
+                    5: approx(-1.69266676902771),
+                    6: approx(-0.898107647895813),
+                    7: approx(-1.725127935409546),
+                    8: approx(0.32057392597198486),
+                },
+                {
+                    0: approx(9.647306442260742),
+                    1: approx(-2.112745523452759),
+                    2: approx(-0.8831668496131897),
+                    3: approx(-2.719135046005249),
+                    4: approx(-0.47473397850990295),
+                    5: approx(-2.2424018383026123),
+                    6: approx(0.6101505160331726),
+                    7: approx(-2.2192084789276123),
+                    8: approx(-0.654518187046051),
+                },
+                {
+                    0: approx(10.364561080932617),
+                    1: approx(-2.240158796310425),
+                    2: approx(-0.9236820340156555),
+                    3: approx(-2.623302936553955),
+                    4: approx(-0.501063346862793),
+                    5: approx(-1.9418426752090454),
+                    6: approx(-0.041013482958078384),
+                    7: approx(-2.12089204788208),
+                    8: approx(-1.2565152645111084),
+                },
+                {
+                    0: approx(9.21961784362793),
+                    1: approx(-2.1359012126922607),
+                    2: approx(-0.1689995676279068),
+                    3: approx(-3.0277929306030273),
+                    4: approx(0.2589835822582245),
+                    5: approx(-2.4426753520965576),
+                    6: approx(0.4815778136253357),
+                    7: approx(-2.3223936557769775),
+                    8: approx(-0.23837946355342865),
+                },
+                {
+                    0: approx(11.023712158203125),
+                    1: approx(-2.0757784843444824),
+                    2: approx(-0.8648976683616638),
+                    3: approx(-2.6088907718658447),
+                    4: approx(-0.781069278717041),
+                    5: approx(-1.7040152549743652),
+                    6: approx(-0.5265535712242126),
+                    7: approx(-1.7141871452331543),
+                    8: approx(-1.0127886533737183),
+                },
+                {
+                    0: approx(9.288057327270508),
+                    1: approx(-1.9795234203338623),
+                    2: approx(0.3095782697200775),
+                    3: approx(-2.781409740447998),
+                    4: approx(-0.4492073655128479),
+                    5: approx(-1.7794851064682007),
+                    6: approx(-0.3115033507347107),
+                    7: approx(-2.0755014419555664),
+                    8: approx(-0.8829556107521057),
+                },
+                {
+                    0: approx(7.968824863433838),
+                    1: approx(-2.0212578773498535),
+                    2: approx(-0.43740716576576233),
+                    3: approx(-3.6762948036193848),
+                    4: approx(-0.012759938836097717),
+                    5: approx(-2.320978879928589),
+                    6: approx(0.6871926784515381),
+                    7: approx(-2.541991949081421),
+                    8: approx(0.32606029510498047),
+                },
+            ]
+        ],
+    ]
+}