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