import gradio as gr
from huggingface_hub import hf_hub_download

from fineTuning import RoBERTaModule

MODEL_REPO_ID = "DornierDo17/MiniRoBERTa_SST2"
WEIGHTS_FILE = "finetuned.pt"

weight_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=WEIGHTS_FILE)


module = RoBERTaModule()
module.load_checkpoint(path=weight_path, location="cpu")


def predict(sentence):
    result = module.inference(sentence)
    return result


gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        label="Enter a piece of the review",
        placeholder="Example: I liked that part of film!",
    ),
    outputs=gr.Textbox(label="Predicted outcome: 0 - negative, 1 - positive "),
    title="Review classfication",
).launch()