mbabanov/devops

1import streamlit as st2import torch3from transformers import AutoTokenizer, AutoModel, pipeline4from torch import nn5 6st.markdown("### Articles classificator.")7 8@st.cache(allow_output_mutation=True)9def get_tokenizer():10    model_name = 'microsoft/deberta-v3-small'11    return AutoTokenizer.from_pretrained(model_name)12 13tokenizer = get_tokenizer()14 15class devops_model(nn.Module):16    def __init__(self):17        super(devops_model, self).__init__()18        self.berta = None19        self.fc = nn.Sequential(20            nn.Linear(768, 768),21            nn.ReLU(),22            nn.Dropout(0.3),23            nn.BatchNorm1d(768),            24            nn.Linear(768, 5),25            nn.LogSoftmax(dim=-1)26        )27        28    def forward(self, train_batch):29        emb = self.berta(**train_batch)['last_hidden_state'].mean(axis=1)30        return self.fc(emb)31 32@st.cache33def LoadModel():34    return torch.load('model_full.pt', map_location=torch.device('cpu'))35 36model = LoadModel()37 38classes = ['Computer Science', 'Mathematics', 'Physics', 'Quantitative Biology', 'Statistics']39 40def process(title, summary):41    text = title + summary42    if not text.strip():43        return ''44    model.eval()45    lines = [text]46    X = tokenizer(lines, padding=True, truncation=True, return_tensors="pt")47    out = model(X)48    probs = torch.exp(out[0])49    sorted_indexes = torch.argsort(probs, descending=True)50    probs_sum = idx = 051    res = []52    while probs_sum < 0.95:53        prob_idx = sorted_indexes[idx]54        prob = probs[prob_idx]55        res.append(f'{classes[prob_idx]}: {prob:.3f}')    56        idx += 157        probs_sum += prob58    return res59    60title = st.text_area("Title", height=30)61 62summary = st.text_area("Summary", height=180)63 64for string in process(title, summary):65    st.markdown(string)