ium_434704/mlflow_prediction_registry.py
Wojciech Jarmosz b3e254c3a5
Some checks failed
s434704-predict-s426206-from-registry/pipeline/head There was a failure building this commit
Add model fetching from mlflow registry
2021-05-23 21:42:55 +02:00

28 lines
555 B
Python

import json
import mlflow
from mlflow.tracking import MlflowClient
import mlflow.pyfunc
import torch
import numpy as np
import pandas as pd
import sys
arguments = sys.argv[1:]
mlflow.set_tracking_uri("http://172.17.0.1:5000")
client = MlflowClient()
version = 1
model_name = "s434704"
input = str(arguments[0])
model = mlflow.pyfunc.load_model(
model_uri=f"models:/{model_name}/{model_version}"
)
with open(f'{model_name}/{input}', 'r') as file:
json_data = json.load(file)
print(model(torch.tensor(np.array(json_data['inputs'])).float()))