import json
import mlflow
import pandas as pd

import sys

arguments = sys.argv[1:]

input = str(arguments[0])

model = mlflow.pytorch.load_model("my_model")

with open('my_model/{input}', 'r') as file:
    json_data = json.load(file)['inputs']

data_into_dict = {idx:x for idx, x in enumerate(json_data)}
print(model.predict(pd.DataFrame([data_into_dict])))