Repository: mlflow/mlflow

Issue:
[BUG] Mlflow transformers fails for fill-mask task for cases where a single data point has multiple masks
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

- Client: 1.x.y
- Tracking server: 1.x.y


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **Python version**:
- **yarn version, if running the dev UI**:


### Describe the problem

For fill mask task, the following scenario is not supported currently:

> Data containing more than one sample where any sample contains more than one mask

[
![Screenshot 2023-10-19 181216](https://github.com/mlflow/mlflow/assets/103428325/a930e08a-dfe2-4cf1-a889-fb79dea5a7e7)
](url)

Example input:
![Screenshot 2023-10-19 181232](https://github.com/mlflow/mlflow/assets/103428325/17736903-b930-484d-98b2-2c98f37c7168)


### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
import pandas as pd
import tempfile
import transformers
import mlflow

architecture = "distilroberta-base"
model = transformers.AutoModelForMaskedLM.from_pretrained(architecture, low_cpu_mem_usage=True)
tokenizer = transformers.AutoTokenizer.from_pretrained(architecture)
pipeline = transformers.pipeline(task="fill-mask", model=model, tokenizer=tokenizer)

model_path = tempfile.tempdir + "/model"
mlflow.transformers.save_model(pipeline, path=model_path)
pyfunc_loaded = mlflow.pyfunc.load_model(model_path)

df = pd.DataFrame({'X': ["I <mask> the whole bowl of <mask>", "I <mask> the whole bowl of <mask>"]})
pyfunc_loaded.predict(df)
```

### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
  File "/opt/miniconda/envs/userenv/lib/python3.8/site-packages/mlflow/pyfunc/__init__.py", line 461, in predict
    return _predict()
  File "/opt/miniconda/envs/userenv/lib/python3.8/site-packages/mlflow/pyfunc/__init__.py", line 447, in _predict
    return self._predict_fn(data, params=params)
  File "/opt/miniconda/envs/userenv/lib/python3.8/site-packages/mlflow/transformers.py", line 1700, in predict
    predictions = self._predict(input_data)
  File "/opt/miniconda/envs/userenv/lib/python3.8/site-packages/mlflow/transformers.py", line 1797, in _predict
    output = self._parse_list_of_multiple_dicts(raw_output, output_key)
  File "/opt/miniconda/envs/userenv/lib/python3.8/site-packages/mlflow/transformers.py", line 2276, in _parse_list_of_multiple_dicts
    return [collection[0][target_dict_key] for collection in output_data]
  File "/opt/miniconda/envs/userenv/lib/python3.8/site-packages/mlflow/transformers.py", line 2276, in <listcomp>
    return [collection[0][target_dict_key] for collection in output_data]
TypeError: list indices must be integers or slices, not str

```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/gateway`: AI Gateway service, Gateway client APIs, third-party Gateway integrations
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/transformers/__init__.py b/mlflow/transformers/__init__.py
--- a/mlflow/transformers/__init__.py
+++ b/mlflow/transformers/__init__.py
@@ -2316,8 +2316,16 @@ def _parse_list_of_multiple_dicts(output_data, target_dict_key):
         Returns the first value of the `target_dict_key` that matches in the first dictionary in a
         list of dictionaries.
         """
+
+        def fetch_target_key_value(data, key):
+            if isinstance(data[0], dict):
+                return data[0][key]
+            return [item[0][key] for item in data]
+
         if isinstance(output_data[0], list):
-            return [collection[0][target_dict_key] for collection in output_data]
+            return [
+                fetch_target_key_value(collection, target_dict_key) for collection in output_data
+            ]
         else:
             return [output_data[0][target_dict_key]]
 


Test patch:
diff --git a/tests/transformers/test_transformers_model_export.py b/tests/transformers/test_transformers_model_export.py
--- a/tests/transformers/test_transformers_model_export.py
+++ b/tests/transformers/test_transformers_model_export.py
@@ -1644,6 +1644,17 @@ def test_fill_mask_pipeline(fill_mask_pipeline, model_path, inference_payload, r
     assert pd_inference == result
 
 
+def test_fill_mask_pipeline_with_multiple_masks(fill_mask_pipeline, model_path):
+    data = ["I <mask> the whole <mask> of <mask>", "I <mask> the whole <mask> of <mask>"]
+
+    mlflow.transformers.save_model(fill_mask_pipeline, path=model_path)
+    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
+
+    inference = pyfunc_loaded.predict(data)
+    assert len(inference) == 2
+    assert all(len(value) == 3 for value in inference)
+
+
 @pytest.mark.parametrize(
     "invalid_data",
     [