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)}") |