wehub-resource-sync
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
import os
|
|
|
|
import mlflow
|
|
from mlflow.models import set_model, set_retriever_schema
|
|
from mlflow.pyfunc import PythonModel
|
|
|
|
test_trace = os.environ.get("TEST_TRACE", "true").lower() == "true"
|
|
|
|
|
|
class MyModel(PythonModel):
|
|
def _call_retriever(self, id):
|
|
return f"Retriever called with ID: {id}. Output: 42."
|
|
|
|
def predict(self, context, model_input):
|
|
return f"Input: {model_input}. {self._call_retriever(model_input)}"
|
|
|
|
def predict_stream(self, context, model_input, params=None):
|
|
yield f"Input: {model_input}. {self._call_retriever(model_input)}"
|
|
|
|
|
|
class MyModelWithTrace(PythonModel):
|
|
def _call_retriever(self, id):
|
|
return f"Retriever called with ID: {id}. Output: 42."
|
|
|
|
@mlflow.trace
|
|
def predict(self, context, model_input):
|
|
return f"Input: {model_input}. {self._call_retriever(model_input)}"
|
|
|
|
@mlflow.trace
|
|
def predict_stream(self, context, model_input, params=None):
|
|
yield f"Input: {model_input}. {self._call_retriever(model_input)}"
|
|
|
|
|
|
model = MyModelWithTrace() if test_trace else MyModel()
|
|
set_model(model)
|
|
set_retriever_schema(
|
|
primary_key="primary-key",
|
|
text_column="text-column",
|
|
doc_uri="doc-uri",
|
|
other_columns=["column1", "column2"],
|
|
)
|