Repository: mlflow/mlflow

Issue:
[BUG] input_example not working properly for pyfunc.PythonModel
### 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)

### Where did you encounter this bug?

Databricks

### Willingness to contribute

No. I cannot contribute a bug fix at this time.

### MLflow version

mlflow==2.8.0

### System information

- Databricks runtime 14.2

### Describe the problem

The use of `input_example` to infer schemas for custom python models, i.e. `pyfunc.PythonModel`, is not working. For example, the example in the [documentation](https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html) does not work. The error message below indicates that the position is not aligned correctly. Context is required to call `predict`. Instead, only `model_input` is given. 



### Tracking information

_No response_

### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from typing import List, Dict
import mlflow

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input: List[str]) -> List[str]:
        return [i.upper() for i in model_input]


mlflow.pyfunc.save_model("model", python_model=MyModel(), input_example=["a"])
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File <command-3509881627683303>, line 10
      6     def predict(self, context, model_input: List[str]) -> List[str]:
      7         return [i.upper() for i in model_input]
---> 10 mlflow.pyfunc.save_model("model", python_model=MyModel(), input_example=["a"])

File /databricks/python/lib/python3.10/site-packages/mlflow/pyfunc/__init__.py:1902, in save_model(path, loader_module, data_path, code_path, conda_env, mlflow_model, python_model, artifacts, signature, input_example, pip_requirements, extra_pip_requirements, metadata, model_config, **kwargs)
   1900 elif isinstance(python_model, PythonModel):
   1901     input_arg_index = 1  # second argument
-> 1902     if signature := _infer_signature_from_type_hints(
   1903         python_model.predict,
   1904         input_arg_index=input_arg_index,
   1905         input_example=input_example,
   1906     ):
   1907         mlflow_model.signature = signature
   1908     elif input_example is not None:

File /databricks/python/lib/python3.10/site-packages/mlflow/models/signature.py:319, in _infer_signature_from_type_hints(func, input_arg_index, input_example)
    317     output_example = func(input_example, params=params) if input_example else None
    318 else:
--> 319     output_example = func(input_example) if input_example else None
    320 output_schema = (
    321     _infer_schema_from_type_hint(hints.output, output_example) if hints.output else None
    322 )
    323 if not any([input_schema, output_schema, params_schema]):

TypeError: MyModel.predict() missing 1 required positional argument: 'model_input'
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `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:
Hi @hoiting Thanks for raising this bug! I can repro it, currently the workaround is to remove the type hint in your `predict` method:
```
from typing import List, Dict
import mlflow

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        # please note that current schema enforcement logic converts List[str] to pandas dataframe
        return model_input

mlflow.pyfunc.save_model("model", python_model=MyModel(), input_example=["a"])
```

I'll raise a PR to fix this.

Resolving patch:
diff --git a/mlflow/models/signature.py b/mlflow/models/signature.py
--- a/mlflow/models/signature.py
+++ b/mlflow/models/signature.py
@@ -313,10 +313,17 @@ def _infer_signature_from_type_hints(func, input_arg_index, input_example=None):
 
     input_schema = _infer_schema_from_type_hint(hints.input, input_example) if hints.input else None
     params_schema = _infer_param_schema(params) if params else None
-    if params and params_key in inspect.signature(func).parameters:
-        output_example = func(input_example, params=params) if input_example else None
+    input_arg_name = _get_arg_names(func)[input_arg_index]
+    if input_example:
+        inputs = {input_arg_name: input_example}
+        if params and params_key in inspect.signature(func).parameters:
+            inputs[params_key] = params
+        # This is for PythonModel's predict function
+        if input_arg_index == 1:
+            inputs["context"] = None
+        output_example = func(**inputs)
     else:
-        output_example = func(input_example) if input_example else None
+        output_example = None
     output_schema = (
         _infer_schema_from_type_hint(hints.output, output_example) if hints.output else None
     )


Test patch:
diff --git a/tests/pyfunc/test_model_export_with_class_and_artifacts.py b/tests/pyfunc/test_model_export_with_class_and_artifacts.py
--- a/tests/pyfunc/test_model_export_with_class_and_artifacts.py
+++ b/tests/pyfunc/test_model_export_with_class_and_artifacts.py
@@ -1464,3 +1464,19 @@ def test_load_model_fails_for_feature_store_models(tmp_path):
         MlflowException, match="mlflow.pyfunc.load_model is not supported for Feature Store models"
     ):
         mlflow.pyfunc.load_model(f"runs:/{run.info.run_id}/model")
+
+
+def test_pyfunc_model_infer_signature_from_type_hints(model_path):
+    class TestModel(mlflow.pyfunc.PythonModel):
+        def predict(self, context, model_input: List[str], params=None) -> List[str]:
+            return model_input
+
+    with mlflow.start_run():
+        model_info = mlflow.pyfunc.log_model(
+            python_model=TestModel(),
+            artifact_path="test_model",
+            input_example=["a"],
+        )
+    pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
+    assert pyfunc_model.metadata.get_input_schema().to_dict() == [{"type": "string"}]
+    pd.testing.assert_frame_equal(pyfunc_model.predict(["a", "b"]), pd.DataFrame(["a", "b"]))