from fastai.text.all import *
import gradio as gr
from blurr.text.modeling.all import *

learn = load_learner("model.pkl")

def predict(inp:str):
    preds = learn.blurr_predict([inp])[0]
    
    preds_dict = dict(zip(preds['class_labels'], preds['probs']))
    preds_dict = sorted(preds_dict.items(), key=operator.itemgetter(1), reverse=True)[:5]
    
    preds_df = pd.DataFrame(preds_dict, columns=['Specialty', 'Probability'])
    preds_df['Probability'] = preds_df['Probability'].apply(lambda x: f"{x*100:.4f}%")
    
    return preds_df
    
intf = gr.Interface(fn=predict, 
                    inputs=gr.inputs.Textbox(label="What are the symptoms?"), 
                    outputs=gr.outputs.Dataframe(),
                    title="Medical Specialty Classification from Symptoms",
                    description="Given a descriptive prompt of symptoms, the model classifies which medical specialty might the symptoms be related too.",
                    examples=["I have been having a headache for two weeks", 
                              "I have rashes on my skin", 
                              "I have been coughing for more than a month"]
)
intf.launch()