|
import streamlit as st |
|
from transformers import DistilBertModel, DistilBertTokenizer |
|
import torch |
|
import matplotlib.image as mpimg |
|
|
|
img = mpimg.imread('./460.jpeg') |
|
model_path = './models/pytorch_distilbert.bin' |
|
vocab_path = './models/vocab_distilbert.bin' |
|
device = torch.device('cpu') |
|
MAX_LEN = 512 |
|
|
|
labels_description = {0: 'Computer Science', |
|
1: 'Economics', |
|
2: 'Electrical Engineering and Systems Science', |
|
3: 'Mathematics', |
|
4: 'Physics', |
|
5: 'Quantitative Biology', |
|
6: 'Quantitative Finance', |
|
7: 'Statistics'} |
|
|
|
|
|
class DistillBERTClass(torch.nn.Module): |
|
def __init__(self): |
|
super(DistillBERTClass, self).__init__() |
|
self.l1 = DistilBertModel.from_pretrained("distilbert-base-cased") |
|
self.pre_classifier = torch.nn.Linear(768, 768) |
|
self.dropout = torch.nn.Dropout(0.3) |
|
self.classifier = torch.nn.Linear(768, 8) |
|
|
|
def forward(self, input_ids, attention_mask): |
|
output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask) |
|
hidden_state = output_1[0] |
|
pooler = hidden_state[:, 0] |
|
pooler = self.pre_classifier(pooler) |
|
pooler = torch.nn.ReLU()(pooler) |
|
pooler = self.dropout(pooler) |
|
output = self.classifier(pooler) |
|
return output |
|
|
|
|
|
def predict(text, model, human_readable=True): |
|
model.eval() |
|
text = " ".join(text.split()) |
|
inputs = tokenizer.encode_plus( |
|
text, |
|
None, |
|
add_special_tokens=True, |
|
max_length=MAX_LEN, |
|
pad_to_max_length=True, |
|
return_token_type_ids=True, |
|
truncation=True |
|
) |
|
ids = torch.tensor(inputs['input_ids'], dtype=torch.long) |
|
ids = torch.reshape(ids, (1, MAX_LEN)) |
|
mask = torch.tensor(inputs['attention_mask'], dtype=torch.long) |
|
mask = torch.reshape(mask, (1, MAX_LEN)) |
|
with torch.no_grad(): |
|
outputs = torch.softmax(model(ids, mask), dim=-1)[0].tolist() |
|
result = [] |
|
for i, v in enumerate(outputs): |
|
result.append((v, i)) |
|
result.sort(reverse=True) |
|
pr = 0.0 |
|
index = 0 |
|
answer = [] |
|
while pr < 0.95: |
|
pr += result[index][0] |
|
if not human_readable: |
|
answer.append(result[index][1]) |
|
else: |
|
answer.append(labels_description[result[index][1]] + " - {:.2f}%".format(100 * result[index][0])) |
|
index += 1 |
|
|
|
return answer |
|
|
|
|
|
tokenizer = DistilBertTokenizer.from_pretrained(vocab_path) |
|
model = torch.load(model_path, map_location=torch.device(device)) |
|
|
|
st.markdown("### Hi! This is a service for determining the subject of an article.") |
|
st.markdown("It can predict the following topics:\n" |
|
"* Computer Science\n" |
|
"* Economics\n" |
|
"* Electrical Engineering and Systems Science\n" |
|
"* Mathematics\n" |
|
"* Physics\n" |
|
"* Quantitative Biology\n" |
|
"* Quantitative Finance\n" |
|
"* Statistics\n") |
|
st.markdown("#### Just write the title and abstract in the areas below and click the \"Analyze\" button.") |
|
|
|
title = st.text_area("Title") |
|
|
|
abstract = st.text_area("Abstract") |
|
|
|
if st.button('Analyse'): |
|
with st.spinner("Wait..."): |
|
if not title and not abstract: |
|
st.error(f"You haven't written anything.") |
|
elif not title: |
|
st.error(f"You haven't written a title.") |
|
else: |
|
pred = predict(title+"\n"+abstract, model.to(device)) |
|
st.success("\n\n".join(pred)) |