import gradio as gr
import re

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

ckpt = "Narrativaai/fake-news-detection-spanish"

tokenizer = AutoTokenizer.from_pretrained(ckpt)

model = AutoModelForSequenceClassification.from_pretrained(ckpt)

classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)


def prediction(header, text):
        results = classifier(header + " [SEP] " + text)
        return results[0]["label"], results[0]["score"]


gradio_ui = gr.Interface(
    fn=prediction,
    title="Fake News Detector (Spanish)",
    description="Type/Paste a post and check if it is Real or Fake",
    inputs=[
        gr.inputs.Textbox(lines=1, label="Type/Paste your headline here"),
        gr.inputs.Textbox(lines=6, label="Type/Paste the article body here"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Label"),
        gr.outputs.Textbox(label="Score"),
    ],
    examples=[
        ["El Real Madrid cerca de la bancarrota", "El Real Madrid cerca de la bancarrota"],
        ["El FC Barcelona está en bancarrota", ""],
    ]
)

gradio_ui.launch()