import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModel, pipeline
from torch import nn

st.markdown("### Articles classificator.")

@st.cache(allow_output_mutation=True)
def get_tokenizer():
    model_name = 'microsoft/deberta-v3-small'
    return AutoTokenizer.from_pretrained(model_name)

tokenizer = get_tokenizer()

class devops_model(nn.Module):
    def __init__(self):
        super(devops_model, self).__init__()
        self.berta = None
        self.fc = nn.Sequential(
            nn.Linear(768, 768),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.BatchNorm1d(768),            
            nn.Linear(768, 5),
            nn.LogSoftmax(dim=-1)
        )
        
    def forward(self, train_batch):
        emb = self.berta(**train_batch)['last_hidden_state'].mean(axis=1)
        return self.fc(emb)

@st.cache
def LoadModel():
    return torch.load('model_full.pt', map_location=torch.device('cpu'))

model = LoadModel()

classes = ['Computer Science', 'Mathematics', 'Physics', 'Quantitative Biology', 'Statistics']

def process(title, summary):
    text = title + summary
    if not text.strip():
        return ''
    model.eval()
    lines = [text]
    X = tokenizer(lines, padding=True, truncation=True, return_tensors="pt")
    out = model(X)
    probs = torch.exp(out[0])
    sorted_indexes = torch.argsort(probs, descending=True)
    probs_sum = idx = 0
    res = []
    while probs_sum < 0.95:
        prob_idx = sorted_indexes[idx]
        prob = probs[prob_idx]
        res.append(f'{classes[prob_idx]}: {prob:.3f}')    
        idx += 1
        probs_sum += prob
    return res
    
title = st.text_area("Title", height=30)

summary = st.text_area("Summary", height=180)

for string in process(title, summary):
    st.markdown(string)