Spaces:
Runtime error
Runtime error
import streamlit as st | |
from transformers import BertTokenizer, BertModel | |
import torch | |
TAGS_CLASSES = ['cs.CV', 'cs.LG', 'cs.AI', 'stat.ML', 'cs.CL', 'cs.NE', 'cs.IR', | |
'math.OC', 'cs.RO', 'cs.LO', 'cs.SI', 'cs.DS', 'cs.IT', 'math.IT', | |
'q-bio.NC', 'stat.ME', 'cs.HC', 'cs.CR', 'cs.DC', 'cs.SD', 'cs.CY', | |
'stat.AP', 'cs.MM', 'math.ST', 'stat.TH', 'cs.DB', 'cs.GT', 'I.2.7', | |
'physics.soc-ph', 'cs.CE', 'cs.SY', 'cs.MA', 'stat.CO', 'cs.NA', | |
'q-bio.QM', 'cs.GR', 'cs.CC', 'physics.data-an', 'cs.SE', 'math.NA', | |
'math.PR', 'quant-ph', 'cs.DL', 'cs.NI', 'I.2.6', 'cs.PL', | |
'cond-mat.dis-nn', 'nlin.AO', 'cmp-lg', 'cs.DM', 'Other'] | |
class BERTClf(torch.nn.Module): | |
def __init__(self): | |
super(BERTClf, self).__init__() | |
self.bert_model = BertModel.from_pretrained('bert-base-uncased', return_dict=True) | |
self.dropout = torch.nn.Dropout(0.1) | |
self.linear = torch.nn.Linear(768, len(TAGS_CLASSES)) | |
self.sigm = nn.Sigmoid() | |
def forward(self, input_ids, attn_mask, token_type_ids): | |
output = self.bert_model( | |
input_ids, | |
attention_mask=attn_mask, | |
token_type_ids=token_type_ids | |
) | |
output_dropout = self.dropout(output.pooler_output) | |
output = self.sigm(self.linear(output_dropout)) | |
return output | |
MAX_LEN = 128 | |
st.markdown("# Paper classification") | |
st.markdown("### Title of paper") | |
# ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter | |
title = st.text_area("TEXT HERE") | |
# ^-- показать текстовое поле. В поле text лежит строка, которая находится там в данный момент | |
st.markdown("### Summary of paper") | |
summary = st.text_area("TEXT HERE", key = "last_name") | |
text = 'Title: ' + title + '\nSummary: ' + summary | |
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') | |
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') | |
model = torch.load('model_5_eps', map_location=device) | |
encodings = tokenizer.encode_plus( | |
text, | |
None, | |
add_special_tokens=True, | |
max_length=MAX_LEN, | |
padding='max_length', | |
return_token_type_ids=True, | |
truncation=True, | |
return_attention_mask=True, | |
return_tensors='pt' | |
) | |
model.eval() | |
with torch.no_grad(): | |
input_ids = encodings['input_ids'].to(device, dtype=torch.long) | |
attention_mask = encodings['attention_mask'].to(device, dtype=torch.long) | |
token_type_ids = encodings['token_type_ids'].to(device, dtype=torch.long) | |
output = model(input_ids, attention_mask, token_type_ids) | |
final_output = output.cpu().detach().numpy().tolist() | |
pred = ([(k,v) for k, v in sorted(zip(TAGS_CLASSES, final_output[0]), key=lambda item: -item[1])])# тут уже знакомый вам код с huggingface.transformers -- его можно заменить на что угодно от fairseq до catboost | |
probs = 0 | |
ans = [] | |
for k, v in pred: | |
if probs > 0.95: | |
break | |
probs += v | |
ans.append(k) | |
st.markdown(f"{', '.join(ans)}") |