import streamlit as st
import torch
import transformers

st.markdown("### Articles classificator.")
# st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)

@st.cache
def LoadModel():
    return torch.load('model.pt'), AutoTokenizer.from_pretrained('bert-base-uncased')()

model, tokenizer = LoadModel()

def process(title, summary):
    text = title + summary
    model.eval()
    lines = [text]
    X = tokenizer(lines, padding=True, truncation=True, return_tensors="pt")
    out = model(X)
    probs = torch.exp(out[0])
    return probs
    
title = st.text_area("Title")

summary = st.text_area("Summary")

st.markdown(f"{process(title, summary)}")