Update app.py

app.py

@@ -23,7 +23,8 @@ def load_model():
     model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=2)
     base_model = SentenceTransformer(base_model_checkpoint)
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, fast=True)
-    return model, base_model, tokenizer
+    data = load_dataset(data_checkpoint, split="train")
+    return model, base_model, tokenizer, data
 
 def sigmoid(x):
     return 1 / (1 + np.exp(-x))
@@ -32,8 +33,7 @@ input_column, reference_column = st.columns(2)
 input_column.write('# Fake News Detection AI')
 
 with st.spinner("Loading Model..."):
-    model, base_model, tokenizer = load_model()
-    data = load_dataset(data_checkpoint, split="train")
+    model, base_model, tokenizer, data = load_model()
 
 user_input = input_column.text_input("Article url")
 submit = input_column.button("submit")
@@ -71,16 +71,16 @@ if submit:
             input_column.text(f"{int(result[prediction]*100)}% confidence")
             input_column.progress(result[prediction])
             
-            title_embeddings = base_model.encode(title)
-            similarity_score = cosine_similarity(
-                [title_embeddings],
-                data["embeddings"]
-            ).flatten()
-            sorted = np.argsort(similarity_score)[::-1].tolist()
-            for i in sorted[:5]:
-                reference_column.write(f"""
-                <a href={data["url"][i]}><small>turnbackhoax.id</small></a>
-                <h5>{data["title"][i]}</h5>
-                """, unsafe_allow_html=True)
-                with reference_column.expander("read content"):
-                    st.write(data["text"][i])
+        title_embeddings = base_model.encode(title)
+        similarity_score = cosine_similarity(
+            [title_embeddings],
+            data["embeddings"]
+        ).flatten()
+        sorted = np.argsort(similarity_score)[::-1].tolist()
+        for i in sorted[:5]:
+            reference_column.write(f"""
+            <a href={data["url"][i]}><small>turnbackhoax.id</small></a>
+            <h5>{data["title"][i]}</h5>
+            """, unsafe_allow_html=True)
+            with reference_column.expander("read content"):
+                st.write(data["text"][i])