import gradio as gr
from huggingface_hub import hf_hub_download
import fasttext, re

REPO = "NbAiLab/nb-nbnn-lid"
FILE = "nb-nbnn-lid.ftz"

MODEL_PATH = hf_hub_download(repo_id=REPO, filename=FILE)
ft = fasttext.load_model(MODEL_PATH)

MAP = {"__label__nob": "nb", "__label__nno": "nn"}  # Bokmål / Nynorsk

def _clean(text: str) -> str:
    text = text or ""
    text = re.sub(r"[\r\n\t]+", " ", text).strip()
    return re.sub(r"\s{2,}", " ", text)

def classify(text: str) -> tuple[str, str]:
    t = _clean(text)
    if not t:
        return "", ""
    label, score = ft.predict(t)  # single text → top-1 label
    predicted_label = MAP.get(label[0], "nb")
    confidence = f"{score[0]:.4f}"
    return predicted_label, confidence

demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="Text"),
    outputs=[
        gr.Textbox(label="Predicted Language (nb or nn)"),
        gr.Textbox(label="Confidence Score")
    ],
    title="Bokmål / Nynorsk classifier",
    description="### Bokmål / Nynorsk classifier",
    api_name="predict"
)

if __name__ == "__main__":
    # turn off SSR to avoid noisy log line; not required
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)