import gradio as gr
import torch

import torch.nn.functional as F
from transformers import BertTokenizer, BertForSequenceClassification


tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("bert-base-uncased")

model.load_state_dict(torch.load('model_after_train.pt', map_location=torch.device('cpu')), strict=False)
model.eval()
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')


def preprocess_text(text):
    parts = []

    text_len = len(text.split(' '))
    delta = 300
    max_parts = 5
    nb_cuts = int(text_len / delta)
    nb_cuts = min(nb_cuts, max_parts)
    
    
    for i in range(nb_cuts + 1):
        text_part = ' '.join(text.split(' ')[i * delta: (i + 1) * delta])
        parts.append(tokenizer.encode(text_part, return_tensors="pt", max_length=500).to(device))

    return parts

def test(text):
    text_parts = preprocess_text(text)
    overall_output = torch.zeros((1,2)).to(device)
    try:
        for part in text_parts:
            if len(part) > 0:
                overall_output += model(part.reshape(1, -1))[0]
    except RuntimeError:
        print("GPU out of memory, skipping this entry.")

    overall_output = F.softmax(overall_output[0], dim=-1)

    value, result = overall_output.max(0)

    term = "fake"
    if result.item() == 0:
        term = "real"

    return term + " at " + str(int(value.item()*100)) + " %"


description = "Fake news detector trained using pre-trained model bert-base-uncased, fine-tuned on https://www.kaggle.com/clmentbisaillon/fake-and-real-news-dataset dataset"
title = "Fake News Detector"

examples = ["CNN - Two people in China are being treated for plague, authorities said Tuesday. It’s the second time the disease, the same one that caused the Black Death, one of the deadliest pandemics in human history, has been detected in the region – in May, a Mongolian couple died from bubonic plague after eating the raw kidney of a marmot, a local folk health remedy."]

iface = gr.Interface(fn=test, inputs="text", outputs="text", title=title,description=description, examples=examples)
iface.launch()