RuinedOustrich/DebertaTextClassify

1#!/usr/bin/env python32# -*- coding: utf-8 -*-3"""4@author: gleb diakonov5"""6 7import streamlit as st8import numpy as np9import torch10import torch.nn as nn11import torch.nn.functional as F12from transformers import AutoTokenizer, AutoModel, AutoConfig13from huggingface_hub import PyTorchModelHubMixin14import re15 16MAX_LEN = 25617 18target_idxs =  ['Astrophysics', 'Condensed Matter', 'Computer Science',19                    'Economics', 'Electrical Engineering and Systems Science',20                    'General Relativity and Quantum Cosmology', 'High Energy Physics - Experiment', 21                    'High Energy Physics - Lattice', 'High Energy Physics - Phenomenology', 22                    'High Energy Physics - Theory', 'Mathematics', 'Mathematical Physics', 23                    'Nonlinear Sciences', 'Nuclear Experiment', 'Nuclear Theory',24                    'Physics', 'Quantitative Biology', 'Quantitative Finance', 'Quantum Physics', 'Statistics']25 26def is_ok(text):27    if not text:28        match = True29    else:30        match = re.match("[a-z]+\d*", text)31                         32    return bool(match)33 34@st.cache_data35def define_tokenizer():36    tokenizer = AutoTokenizer.from_pretrained('./token')37    return tokenizer38 39def preprocess(text):40    41    tokenizer = define_tokenizer()42    43    encoded_text = tokenizer.encode_plus(44        text,45        max_length=MAX_LEN,46        add_special_tokens=True,47        return_token_type_ids=False,48        padding='max_length',49        truncation = True,50        return_attention_mask=True,51        return_tensors='pt',52        )53 54    input_ids = encoded_text['input_ids']55    attention_mask = encoded_text['attention_mask']56    57    return input_ids, attention_mask58 59class MeanPooling(nn.Module):60    def __init__(self):61        super(MeanPooling, self).__init__()62        63    def forward(self, last_hidden_state, attention_mask):64        input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()65        sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)66        sum_mask = input_mask_expanded.sum(1)67        sum_mask = torch.clamp(sum_mask, min=1e-9)68        mean_embeddings = sum_embeddings / sum_mask69        return mean_embeddings70 71class BERTClass(torch.nn.Module, PyTorchModelHubMixin):72    def __init__(self, num_classes = 20):73        super(BERTClass, self).__init__()74        config = AutoConfig.from_pretrained("./bert/config.json")75        self.bert_model = AutoModel.from_config(config)76        self.dropout = torch.nn.Dropout(0.3)77        self.batchnorm = nn.BatchNorm1d(768)78        self.pooler = MeanPooling()79        self.linear = torch.nn.Linear(768, num_classes)80 81    def forward(self, input_ids, attn_mask):82        output = self.bert_model(83            input_ids, 84            attention_mask=attn_mask,85        )86        output = self.pooler(output.last_hidden_state, attn_mask)87        output_dropout = self.dropout(output)88        output = self.linear(output_dropout)89        return output90 91@st.cache_data92def configure_model():93    model = BERTClass.from_pretrained("./deberta-arxiv-model")94    return model95 96def predict(text):97    98    model = configure_model()99    100    model.eval()101    102    with torch.no_grad():103    104        input_ids, attention_mask = preprocess(text)105        preds = model(input_ids, attention_mask)106        output = F.softmax(preds, dim = 1).detach()107        output = output.flatten().numpy()108        output = {tag: round(float(prob)*100, 2) for tag, prob in zip(target_idxs, output)}109        outputs = {k: v for k, v in sorted(output.items(), reverse = True, key = lambda x: x[1])}110    return outputs111 112 113if __name__ == '__main__':114    115    st.markdown("<h1 style='text-align: center;'>Find out the topic of the article</h1>", unsafe_allow_html=True)116    st.markdown("<h2 style='text-align: center;'>Please enter title or summary</h2>", unsafe_allow_html=True)117    st.markdown("<h4 style='text-align: center;'>(enter both for better result)</h4>", unsafe_allow_html=True)118    form = st.form("my_form")119    title = form.text_input("TITLE")120    summary = form.text_area("SUMMARY")121    button = form.form_submit_button("Submit")122    summary = summary.lower()123    title = title.lower()124    if button:125        if not title and not summary:126            st.write("**PLEASE ENTER SOMETHING!**")127        else:128            text = title + ". " + summary129            if not is_ok(summary) and is_ok(title):130                st.write("**INCORRECT INPUT FORMAT: SUMMARY**")131            if is_ok(summary) and not is_ok(title):132                st.write("**INCORRECT INPUT FORMAT: TITLE**")133            elif not is_ok(title) and not is_ok(summary):134                st.write("**INCORRECT INPUT FORMAT: TITLE, SUMMARY**")135            elif len(summary.split()) in (1,2,3) and not title:136                st.write("**There are too few words in summary, result can be bad. Make shure you enter full text**")137            elif len(title.split()) in (1,2,3) and not summary:138                st.write("**There are too few words in title, result can be bad. Make shure you enter full text**")139            elif len(title.split()) in (1,2,3) and len(summary.split()) == 1:140                st.write("**There are too few words in title and summary, result can be bad. Make shure you enter full text**")141            else:142                outputs = predict(text)143                sums_probs = []144                for k, v in outputs.items(): 145                    st.write(f'**:blue[{k}]** - **:green[{v}%]**')146                    sums_probs.append(v)147                    if sum(sums_probs) >= 95:148                        break149                150                151