cleaner code

app.py

@@ -6,10 +6,10 @@ import time
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer
 from Scraper import Scrap
 
+st.set_page_config(layout="wide")
 
 model_checkpoint = "Rifky/FND"
-label = {0: "Fakta", 1: "Hoax"}
-
+label = {0: "valid", 1: "fake"}
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 def load_model():
@@ -17,23 +17,21 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, fast=True)
     return Trainer(model=model), tokenizer
 
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
 
-st.write('# Fake News Detection AI')
+input_column, reference_column = st.columns(2, gap="medium")
+input_column.write('# Fake News Detection AI')
 
 with st.spinner("Loading Model..."):
     model, tokenizer = load_model()
 
-user_input = st.text_area("Put article url or the full text", help="the text you want to analyze", height=200)
-submit = st.button("submit")
+user_input = input_column.text_input("Article url")
+submit = input_column.button("submit")
 
-def sigmoid(x):
-    return 1 / (1 + np.exp(-x))
 
 if submit:
     last_time = time.time()
-
-    text = ""
-
     with st.spinner("Reading Article..."):
         if user_input:
             if user_input[:4] == 'http':
@@ -45,33 +43,24 @@ if submit:
         text = re.sub(r'\n', ' ', text)
 
         with st.spinner("Computing..."):
-            text_len = len(text.split(" "))
-            if text_len > 512:
-                texts = []
-                for i in range(text_len // 512):
-                    texts.append(" ".join(text.split(" ")[i * 512:(i + 1) * 512]))
-
-                texts.append(" ".join(text.split(" ")[(text_len // 512) + 1:text_len % 512]))
-
-                for i in range(len(texts)):
-                    texts[i] = tokenizer(texts[i], max_length=512, truncation=True, padding="max_length")
-
-                results = model.predict(texts)[0]
-                result = [0, 0]
-                for i in range(len(results)):
-                    result[0] += sigmoid(results[i][0])
-                    result[1] += sigmoid(results[i][1])
-
-                result[0] /= len(results)
-                result[1] /= len(results)
-
-            else:
-                text = tokenizer(text, max_length=512, truncation=True, padding="max_length")
-                result = model.predict([text])[0][0]
+            text = text.split()
+            text_len = len(text)
+
+            sequences = []
+            for i in range(text_len // 512):
+                sequences.append(" ".join(text[i * 512: (i + 1) * 512]))
+            sequences.append(" ".join(text[text_len - (text_len % 512) : text_len]))
+            sequences = [tokenizer(i, max_length=512, truncation=True, padding="max_length") for i in sequences]
+
+            predictions = model.predict(sequences)[0]
+            result = [
+                np.sum([sigmoid(i[0]) for i in predictions]) / len(predictions), 
+                np.sum([sigmoid(i[1]) for i in predictions]) / len(predictions)
+            ]
                 
             print (f'\nresult: {result}')
-
-            st.markdown(f"<small>Compute Finished in {int(time.time() - last_time)} seconds</small>", unsafe_allow_html=True)
-
+            input_column.markdown(f"<small>Compute Finished in {int(time.time() - last_time)} seconds</small>", unsafe_allow_html=True)
             prediction = np.argmax(result, axis=-1)
-            st.success(f"Prediction: {label[prediction]}")
+            input_column.success(f"This news is {label[prediction]}.")
+            st.text(f"{int(result[prediction]*100)}% confidence")
+            input_column.progress(result[prediction])