File size: 2,177 Bytes
dcd4afb 3479b02 dcd4afb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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() |