|
|
import streamlit as st |
|
|
import torch |
|
|
import transformers |
|
|
|
|
|
st.markdown("### Articles classificator.") |
|
|
|
|
|
|
|
|
@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)}") |
