app.py

import streamlit as st

from transformers import DistilBertTokenizer, DistilBertModel
import torch
import torch.nn as nn

from utils import get_answer_with_desc

MY_LINEAR_NAME = "my_linear_logits_3"

st.markdown("## Classifying articles on computer science!")
st.markdown("<img width=400px src='https://img.freepik.com/free-photo/young-pretty-student-overwhelmed-with-books_272645-183.jpg?size=626&ext=jpg'>", unsafe_allow_html=True)

title = st.text_area("Enter the title of the article")
abstract = st.text_area("Enter the abstract of the article")
text = title + " " + abstract

tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-cased")
model = DistilBertModel.from_pretrained("distilbert-base-cased")

n_classes = 40
my_linear = nn.Linear(in_features=768, out_features=n_classes, bias=True)
my_linear.load_state_dict(torch.load(MY_LINEAR_NAME))

st.markdown(get_answer_with_desc(text, model, tokenizer, my_linear))