Spaces:
Runtime error
Runtime error
File size: 4,563 Bytes
f640140 bf64a07 f640140 5ebbffe b4a581e 606c75b 5ebbffe f640140 |
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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import os
import json
import numpy as np
import pandas as pd
import streamlit as st
import torch
import torch.nn.functional as F
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
@st.cache(allow_output_mutation=True)
def init_model():
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased')
model = DistilBertForSequenceClassification.from_pretrained('khizon/distilbert-unreliable-news-eng-4L', num_labels = 2)
return tokenizer, model
def download_dataset():
url = 'https://drive.google.com/drive/folders/11mRvsHAkggFEJvG4axH4mmWI6FHMQp7X?usp=sharing'
data = 'data/nela_gt_2018_site_split'
os.system(f'gdown --folder {url} -O {data}')
@st.cache(allow_output_mutation=True)
def jsonl_to_df(file_path):
with open(file_path) as f:
lines = f.read().splitlines()
df_inter = pd.DataFrame(lines)
df_inter.columns = ['json_element']
df_inter['json_element'].apply(json.loads)
return pd.json_normalize(df_inter['json_element'].apply(json.loads))
@st.cache
def load_test_df():
file_path = os.path.join('data', 'nela_gt_2018_site_split', 'test.jsonl')
test_df = jsonl_to_df(file_path)
test_df = pd.get_dummies(test_df, columns = ['label'])
return test_df
@st.cache(allow_output_mutation=True)
def predict(model, tokenizer, data):
labels = data[['label_0', 'label_1']]
labels = torch.tensor(labels, dtype=torch.float32)
encoding = tokenizer.encode_plus(
data['title'],
' [SEP] ' + data['content'],
add_special_tokens=True,
max_length = 512,
return_token_type_ids = False,
padding = 'max_length',
truncation = 'only_second',
return_attention_mask = True,
return_tensors = 'pt'
)
output = model(**encoding)
return correct_preds(output['logits'], labels)
@st.cache(allow_output_mutation=True)
def predict_new(model, tokenizer, title, content):
encoding = tokenizer.encode_plus(
title,
' [SEP] ' + content,
add_special_tokens=True,
max_length = 512,
return_token_type_ids = False,
padding = 'max_length',
truncation = 'only_second',
return_attention_mask = True,
return_tensors = 'pt'
)
output = model(**encoding)
preds = F.softmax(output['logits'], dim = 1)
p_idx = torch.argmax(preds, dim = 1)
return 'reliable' if p_idx > 0 else 'unreliable'
def correct_preds(preds, labels):
preds = torch.nn.functional.softmax(preds, dim = 1)
p_idx = torch.argmax(preds, dim=1)
l_idx = torch.argmax(labels, dim=0)
pred_label = 'reliable' if p_idx > 0 else 'unreliable'
correct = True if (p_idx == l_idx).sum().item() > 0 else False
return pred_label, correct
if __name__ == '__main__':
if not os.path.exists('data/nela_gt_2018_site_split/test.jsonl'):
download_dataset()
df = load_test_df()
tokenizer, model = init_model()
st.title("Unreliable News classifier")
mode = st.radio(
'', ('Test article', 'Input own article')
)
if mode == 'Test article':
if st.button('Get random article'):
idx = np.random.randint(0, len(df))
sample = df.iloc[idx]
prediction, correct = predict(model, tokenizer, sample)
label = 'reliable' if sample['label_1'] > sample['label_0'] else 'unreliable'
st.header(sample['title'])
if correct:
st.success(f'Prediction: {prediction}')
else:
st.error(f'Prediction: {prediction}')
st.caption(f'Source: {sample["source"]} ({label})')
# if len(sample['content']) > 300:
# sample['content'] = sample['content'][:300]
temp = []
for idx, word in enumerate(sample['content'].split()):
if (idx%5 == 0) and idx > 0:
word = '▒'*len(word)
temp.append(word)
sample['content'] = ' '.join(temp)
st.markdown(sample['content'])
else:
title = st.text_input('Article title', 'Test title')
content = st.text_area('Article content', 'Lorem ipsum')
if st.button('Submit'):
pred = predict_new(model, tokenizer, title, content)
st.markdown(f'Prediction: {pred}')
# st.success('success') |