Как передать данные из конечной точки обслуживания модели Databricks?

Я разрабатываю приложение, использующее Custom PyFunc из MLflow. Приложение включает в себя несколько компонентов, таких как перефразирование запроса, обнаружение намерений, извлечение фрагментов и генерация ответов. Весь процесс занимает около 7–8 секунд, причем первые 3–4 секунды используются для извлечения фрагмента, а оставшиеся 3–4 секунды – для генерации ответа.

Чтобы улучшить взаимодействие с пользователем, я хочу начать потоковую передачу окончательного ответа от OpenAI с 4-й секунды, чтобы пользователю не приходилось ждать полные 7-8 секунд. Однако я столкнулся с проблемой, из-за которой я не могу вернуть генератор из обычного метода прогнозирования в MLflow, получая следующую ошибку:

mlflow.exceptions.MlflowException: Encountered an unexpected error while converting model response to JSON. Error: 'Object of type generator is not JSON serializable.'

Я обнаружил, что Predict_stream в документации MLflow может возвращать генератор. Это хорошо работает с load_model в записной книжке, но когда я обновляю модель в конечной точке обслуживания модели, я получаю ответ None.

Вот собственный пример, который я нашел в документации MLflow:

import mlflow

# Define a custom model that supports streaming
class StreamableModel(mlflow.pyfunc.PythonModel):
    def predict_stream(self, context, model_input, params=None):
        # Yielding elements one at a time
        for element in ["a", "b", "c", "d", "e"]:
            yield element

Вот код для регистрации модели:

with mlflow.start_run():
    # Log the custom Python model
    mlflow.pyfunc.log_model("stream_model", python_model=StreamableModel(),registered_model_name = "stream_model",streamable=True)

Есть ли способ вернуть генератор из конечной точки обслуживания модели?