Upload app.py

app.py

@@ -0,0 +1,157 @@
+import streamlit as st
+import time
+
+# model part
+
+import json
+import torch
+from torch import nn
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+with open('categories_with_names.json', 'r') as f:
+    cat_with_names = json.load(f)
+with open('categories_from_model.json', 'r') as f:
+    categories_from_model = json.load(f)
+
+@st.cache_resource
+def load_models_and_tokenizer():
+    tokenizer = AutoTokenizer.from_pretrained("oracat/bert-paper-classifier-arxiv")
+
+    model_titles = AutoModelForSequenceClassification.from_pretrained(
+        "powerful_model_titles/checkpoint-13472", num_labels=len(categories_from_model), problem_type="multi_label_classification"
+    )
+    model_titles.eval()
+    model_abstracts = AutoModelForSequenceClassification.from_pretrained(
+        "powerful_model_abstracts/checkpoint-13472", num_labels=len(categories_from_model), problem_type="multi_label_classification"
+    )
+    model_abstracts.eval()
+    
+    return model_titles, model_abstracts, tokenizer
+
+model_titles, model_abstracts, tokenizer = load_models_and_tokenizer()
+
+def categorize_text(title: str | None = None, abstract: str | None = None, progress_bar = None):
+    if title is None and abstract is None:
+        raise ValueError('title is None and abstract is None')
+
+    models_to_run = 2 if (title is not None and abstract is not None) else 1
+        
+    proba_title = None
+    if title is not None:
+        progresses = (10, 30) if models_to_run == 2 else (20, 60)
+        my_bar.progress(progresses[0], text='computing titles')
+        input_tok = tokenizer(title, return_tensors='pt')
+        with torch.no_grad():
+            logits = model_titles(**input_tok)['logits']
+            proba_title = torch.sigmoid(logits)[0]
+        my_bar.progress(progresses[1], text='computed titles')
+            
+    proba_abstract = None
+    if abstract is not None:
+        progresses = (40, 70) if models_to_run == 2 else (20, 60)
+        my_bar.progress(progresses[0], text='computing abstracts')
+        input_tok = tokenizer(abstract, return_tensors='pt')
+        with torch.no_grad():
+            logits = model_abstracts(**input_tok)['logits']
+            proba_abstract = torch.sigmoid(logits)[0]
+        my_bar.progress(progresses[0], text='computed abstracts')
+            
+    if title is None:
+        proba = proba_abstract
+    elif abstract is None:
+        proba = proba_title
+    else:
+        proba = proba_title * 0.1 + proba_abstract * 0.9
+
+    progresses = (80, 90) if models_to_run == 2 else (70, 90)
+
+    my_bar.progress(progresses[0], text='computed proba')
+    
+    sorted_proba, indices = torch.sort(proba, descending=True)
+    my_bar.progress(progresses[1], text='sorted proba')
+    to_take = 1
+    while sorted_proba[:to_take].sum() < 0.95 and to_take < len(categories_from_model):
+        to_take += 1
+    output = [(cat_with_names[categories_from_model[index]], proba[index].item())
+            for index in indices[:to_take]]
+    my_bar.progress(100, text='generated output')
+    return output
+
+# front part
+
+st.markdown("<h1 style='text-align: center;'>Classify your paper!</h1>", unsafe_allow_html=True)
+
+
+if "title" not in st.session_state:
+    st.session_state.title = ""
+if "abstract" not in st.session_state:
+    st.session_state.abstract = ""
+if "title_input_key" not in st.session_state:
+    st.session_state.title_input_key = ""
+if "abstract_input_key" not in st.session_state:
+    st.session_state.abstract_input_key = ""
+if "model_type" not in st.session_state:
+    st.session_state.model_type = []
+
+def input_error():
+    if not st.session_state.model_type:
+        return 'you have to select title or abstract'
+    if 'Title' in model_type and not st.session_state.title:
+        return 'Title is empty'
+    if 'Abstract' in model_type and not st.session_state.abstract:
+        return 'Abstract is empty'
+    return ''
+
+
+def clear_input():
+    st.session_state.title = title.title()
+    st.session_state.abstract = abstract.title()
+    if not input_error():
+        if "Title" in st.session_state.model_type:
+           st.session_state.title_input_key = ""
+        if "Abstract" in st.session_state.model_type:
+           st.session_state.abstract_input_key = ""
+
+
+title = st.text_input(r"$\textsf{\Large Title}$", key="title_input_key")
+
+abstract = st.text_input(r"$\textsf{\Large Abstract}$", key="abstract_input_key")
+
+
+model_type = st.multiselect(
+    r"$\textsf{\large Classify by:}$",
+    ['Title', 'Abstract'],
+)
+
+st.session_state.model_type = model_type
+
+if(st.button('Submit', on_click=clear_input)):
+    if input_error():
+        st.error(input_error())
+    else:
+        send_time = time.localtime(time.time())
+        #st.success(f"Submitted {(' and '.join(st.session_state.model_type)).lower()} on {time.strftime('%d.%m %H:%M:%S', send_time)}")
+        model_input = dict()
+        if 'Title' in st.session_state.model_type:
+            model_input['title'] = st.session_state.title
+        if 'Abstract' in st.session_state.model_type:
+            model_input['abstract'] = st.session_state.abstract
+        #st.success(f'{model_input=}')
+        my_bar = st.progress(0, text='starting model')
+        model_result = categorize_text(**model_input, progress_bar=my_bar)
+        st.markdown("<h1 style='text-align: center;'>Classification completed!</h1>", unsafe_allow_html=True)
+        small_categories = []
+        cat, proba = model_result[0]
+        st.write(r"$\textsf{\Large " + f'{cat}: {round(100*proba)}' + r"\%}$")
+        for cat, proba in model_result[1:]:
+            if proba < 0.1:
+                small_categories.append(f'{cat}: {round(100*proba, 1)}' + r"\%")
+            else:
+                st.write(r"$\textsf{\large " + f'{cat}: {round(100*proba)}' + r"\%}$")
+        if small_categories:
+            st.write(', '.join(small_categories))
+
+
+
+
+