works

.gitattributes

@@ -31,3 +31,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,13 +1,52 @@
 import streamlit as st
+import torch
+from torch.nn import functional as F
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import json
+import streamlit.components.v1 as components
 
-st.markdown("### Hello, world!")
-st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
 
-from transformers import pipeline
+if __name__ == '__main__':
+    st.markdown("### Arxiv paper classifier (No guarantees provided)")
 
-pipe = pipeline("ner", "Davlan/distilbert-base-multilingual-cased-ner-hrl")
+    col1, col2 = st.columns([1, 1])
+    col1.image('imgs/akinator_ready.png', width=200)
+    btn = col2.button('Classify!')
 
-text = st.text_area("TEXT HERE")
+    model = AutoModelForSequenceClassification.from_pretrained('checkpoint-3000')
+    tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
 
-st.markdown(f"{pipe(text)}")
+    with open('checkpoint-3000/config.json', 'r') as f:
+        id2label = json.load(f)['id2label']
 
+    id2label = {int(key): value for key, value in id2label.items()}
+    title = st.text_area(label='Input title...', placeholder='Input title...', label_visibility='hidden', height=3)
+    abstract = st.text_area(label='Input title...', placeholder='Input abstract...', label_visibility='hidden', height=10)
+    text = '\n'.join([title, abstract])
+
+    if btn and len(text) == 1:
+        st.error('Title and abstract are empty!')
+
+    if btn and len(text) > 1:
+        tokenized = tokenizer(text)
+
+        with torch.no_grad():
+            out = model(torch.tensor(tokenized['input_ids']).unsqueeze(dim=0))
+        _, ids = torch.sort(-out['logits'])
+        probs = F.softmax(out['logits'][0, ids], dim=1)
+        ids, probs = ids[0], probs[0]
+
+        ptotal = 0
+        result = []
+        for i, prob in enumerate(probs):
+            ptotal += prob
+            result.append(f'{id2label[ids[i].item()]} (prob = {prob.item()})')
+        output = '<br>'.join(result)
+
+        components.html(f'<div>'
+                        f'<div style="height:120px;width:680px;'
+                        f'border:1px solid #ccc;border-color: red;'
+                        f'font:16px/26px Georgia, Garamond, Serif;'
+                        f'overflow:scroll;'
+                        f'color:white;">'
+                        f'{output}</div>')

requirements.txt

@@ -1,3 +1,4 @@
-transformers
-torch
+transformers==4.15.0
+torch==1.12.1
+