print('INFO: import modules')
import gradio as gr
import pickle
from required_classes import *


print('INFO: loading model')
try:
    with open('pretrain_logistic_regression_model.pkl', 'rb') as f:
        model = pickle.load(f)
    model.batch_size = 1
    print('INFO: model loaded')
except Exception as e:
    print(f"ERROR: loading model failed with: {str(e)}")

def classify(text):
    pred_classes = model.predict([text])
    output_text = ' '.join(pred_classes)
    return output_text

print('INFO: starting gradio interface')
iface = gr.Interface(
    enable_queue=True,
    title="ICD10-codes classification",
    description="",
    fn=classify,
    inputs=[gr.inputs.Textbox(label="Input text")],
    outputs=gr.outputs.Textbox(label="Result class"),
)
iface.launch()