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(news_headline):
    learner = from_pretrained_fastai("rajeshradhakrishnan/ml-news-classify-fastai")

    probabilities = learner.predict(news_headline)

    return {LABELS[i]: probabilities[0]['probs'][i] for i in range(len(LABELS))}

interface = gr.Interface(
    predict, 
    inputs="textbox", 
    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 appropriately!", 
    article = "<p style='text-align: center'>Malayalam News Classifier | Demo Model</p>",
    examples=[["ഓഹരി വിപണി തകരുമ്പോള്‍ നിക്ഷേപം എങ്ങനെ സുരക്ഷിതമാക്കാം"], ["വാര്‍ണറുടെ ഒറ്റക്കയ്യന്‍ ക്യാച്ചില്‍ അമ്പരന്ന് ക്രിക്കറ്റ് ലോകം"]],
    # live=True,
    share=True)
interface.launch(debug=True)