Nothing

app.py

@@ -1,12 +1,52 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import DistilBertModel, DistilBertTokenizer
+import torch
 
-st.markdown("## Movie genre classification")
-st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
+model_path = './models/pytorch_distilbert.bin'
+vocab_path = './models/vocab_distilbert.bin'
+device = torch.device('cpu')
+MAX_LEN = 512
 
 
-pipe = pipeline("text-classification", "Tejas3/distillbert_110_uncased_movie_genre")
+def get_labels(text, model, tokenizer, count_labels=8):
+    tokens = tokenizer(text, return_tensors='pt')
+    outputs = model(**tokens)
+    probs = torch.nn.Softmax()(outputs.logits)
 
-text = st.text_area("TEXT HERE")
+    labels = ['Computer_science', 'Economics',
+              'Electrical_Engineering_and_Systems_Science', 'Mathematics',
+              'Physics', 'Quantitative_Biology', 'Quantitative_Finance',
+              'Statistics']
 
-st.markdown(f"{pipe(text)}")
+    sort_lst = sorted([(prob, label) for prob, label in zip(probs.detach().numpy()[0], labels)], key=lambda x: -x[0])
+    cumsum = 0
+    result_labels = []
+    for pair in sort_lst:
+        cumsum += pair[0]
+        if cumsum > 0.95 and len(result_labels) >= 1:
+            return result_labels
+        result_labels.append(pair[1])
+
+
+@st.cache(allow_output_mutation=True)
+def load_model():
+    tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-cased")
+    model = DistilBertModel.from_pretrained("distilbert-base-cased", num_labels=8)
+    model.load_state_dict(torch.load('weight_model'))
+    return model, tokenizer
+
+
+tokenizer = DistilBertTokenizer.from_pretrained(vocab_path)
+model = torch.load(model_path, map_location=torch.device(device))
+
+st.markdown("### Movie genre classification")
+
+text = st.text_area("Write some movie description")
+
+if st.button('Predict'):
+    with st.spinner("Wait..."):
+        if not text:
+            st.error("Write something.")
+        else:
+            pred = predict(text, model.to(device))
+            st.success("\n\n".join(pred))