app.py

import gradio as gr
from transformers import AutoTokenizer
import torch
from fastai.text.all import *
from blurr.text.data.all import *
from blurr.text.modeling.all import *

# Define the path to your model and dataloaders
model_path = "origin-classifier-stage-2.pkl"
dls_path = "dls_origin-classifier_v1.pkl"

learner_inf = load_learner(model_path)
dls = torch.load(dls_path)

class_label_mapping = {label: idx for idx, label in enumerate(learner_inf.dls.vocab)}

def predict_text(text):
    prediction = learner_inf.blurr_predict(text)[0]
    predicted_class_index = prediction['class_index']
    predicted_class_label = list(class_label_mapping.keys())[list(class_label_mapping.values()).index(predicted_class_index)]
    predicted_class_probability = prediction['probs'][predicted_class_index]
    return {predicted_class_label: float(predicted_class_probability)}

iface = gr.Interface(
    fn=predict_text,
    inputs="text",
    outputs=gr.outputs.KeyValues(label="Predicted Class and Probability"),
    title="Food Origin Classification App",
    description="Enter a Recipe, and it will predict the class label.",
)

iface.launch()