import fasttext
from huggingface_hub import hf_hub_download
import gradio as gr
import numpy as np

model_path = hf_hub_download(repo_id="facebook/fasttext-language-identification", filename="model.bin")
model = fasttext.load_model(model_path)

def predict(text, top):
    labels, probabilities = model.predict(text, k=top)
    cleaned_labels = [label.replace('__label__', '') for label in labels]
    result = dict(zip(cleaned_labels, np.array(probabilities)))
    #result = sorted(result, key=lambda x: x[1], reverse=True)
    return result

demo = gr.Interface(
    fn=predict,
    inputs=[
      gr.Textbox(lines=1, placeholder="Text", label="Content"),
      gr.Number(value=5, info='number of predictions that should be returned', minimum=1, maximum=100,  label="Top"),
    ],
    title="Language Identification Demo",
    flagging_mode="never",
    outputs=gr.Label(label="Result"))
demo.launch(share=True, show_api=True)