from pathlib import Path
import gradio as gr
from huggingface_hub import from_pretrained_fastai


LABELS = Path('class_names.txt').read_text().splitlines()



def predict(im):
    learner = from_pretrained_fastai("rajeshradhakrishnan/ml-news-classify-fastai")

    probabilities = learner.predict(im)

    values, indices = torch.topk(probabilities, 5)

    return {LABELS[i]: v.item() for i, v in zip(indices, values)}

interface = gr.Interface(
    predict, 
    inputs="newsheadlines", 
    outputs='label', 
    theme="huggingface", 
    title="Malayalam News Classifier", 
    description="Try to classify news in മലയാളം? Input a few malayalam news headlines and verify whether the model categorized it apporpirately!", 
    article = "<p style='text-align: center'>Malayalam News Classifier | Demo Model</p>",
    live=True)
interface.launch(debug=True)