from huggingface_hub import from_pretrained_fastai
import gradio as gr

repo_id = "pamunarr/P7EjOpc1-LSTM"

learner = from_pretrained_fastai(repo_id)
labels = ["World" , "Nigeria" , "Health" , "Africa" , "Politics"]

def predict(text):
    _ , _ , probs = learner.predict(text)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

gr.Interface(fn=predict, inputs="text", outputs=gr.components.Label(num_top_classes=5)).launch(share=False)