kol/Text_Classification

1import streamlit as st2try:3    import transformers4    from transformers import AutoTokenizer, AutoModel5    import torch6    import torch.nn as nn7    import processing8except BaseException:9    st.write("ERROR: error in importing libs.")10    exit(1)11 12@st.cache13def GetModels():14    model = AutoModel.from_pretrained("./transformer")15    model_head = nn.Sequential(nn.Linear(768, 300), nn.ReLU(), nn.Linear(300, 6))16    model_head.load_state_dict(torch.load("./model.txt", map_location=torch.device('cpu')))17    return model, model_head18 19try:20    model, model_head = GetModels()21except BaseException:22    st.write("ERROR: error in loading model.")23    exit(2)24 25st.header("Text Classification")26st.subheader("Description")27st.write("This app classifies text by title and summary(optional) into 6 classes: Computer Science, Math, Economy and Finance, Statistics, Physics, Biology.")28st.subheader("Classification")29st.write("Enter title and summary (optional)")30 31try:32    title = st.text_area("Title HERE:")33    summary = st.text_area("Summary HERE:")34except BaseException:35    st.write("ERROR: error in input text.")36    exit(3)37    38button = st.button("Get classes")39if button:40    try:41        tokenizer = AutoTokenizer.from_pretrained("./transformer")42    except BaseException:43	    st.write("ERROR: error in loading tokenizer.")44	    exit(4)45	    46    if title == "" and summary == "":47        st.write("Title or summary is necessary for classification.")48    else:49        try:50            classes, probs = processing.MakePrediction(model, model_head, tokenizer, title, summary)51        except BaseException:52            st.write("ERROR: error in classification")53            exit(5)54            55        col1, buff, col2 = st.columns([2, 0.5, 4])56        col1.text("CLASSES:")57        col2.text("PROBABILITIES:")58        for idx, i in enumerate(classes):59            col1.write(i + ":")60            col2.write(probs[idx])61