import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np

model_name = "dancrvlh/Language"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

mapping = {
2:"English",
3:"French",
4:"Portugeese",
5:"Russian",
6:"Sweedish",
7:"Sweedish",
}

def predict(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    predictions = outputs.logits
    return mapping[round(np.exp(predictions.item()))]
    

iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    layout="vertical",
    title="Language Detection",
    description="This model can detect the language your text for English, French, Portugeese, Russian, Sweedish and Unknow",
)

iface.launch(share=True)