ium_434704/mlflow_prediction_registry.py

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()
model_version = 1
model_name = "s426206"

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()))