Reference Searcher

app.py

@@ -6,19 +6,24 @@ import time
 import os
 
 from transformers import AutoModelForSequenceClassification, AutoModel, AutoTokenizer
+from sklearn.metrics.pairwise import cosine_similarity
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
 from Scraper import Scrap
 
 st.set_page_config(layout="wide")
 
 model_checkpoint = "Rifky/FND"
+base_model_checkpoint = "indobenchmark/indobert-base-p1"
 data_checkpoint = "Rifky/turnbackhoax-encoded"
 label = {0: "valid", 1: "fake"}
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
 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, tokenizer
+    return model, base_model, tokenizer
 
 def sigmoid(x):
     return 1 / (1 + np.exp(-x))
@@ -27,7 +32,8 @@ input_column, reference_column = st.columns(2)
 input_column.write('# Fake News Detection AI')
 
 with st.spinner("Loading Model..."):
-    model, tokenizer = load_model()
+    model, base_model, tokenizer = load_model()
+    data = load_dataset(data_checkpoint, split="train")
 
 user_input = input_column.text_input("Article url")
 submit = input_column.button("submit")
@@ -36,11 +42,7 @@ submit = input_column.button("submit")
 if submit:
     last_time = time.time()
     with st.spinner("Reading Article..."):
-        if user_input:
-            if user_input[:4] == 'http':
-                text = Scrap(user_input)
-            else:
-                text = user_input
+        title, text = Scrap(user_input)
 
     if text:
         text = re.sub(r'\n', ' ', text)
@@ -66,4 +68,20 @@ if submit:
             prediction = np.argmax(result, axis=-1)
             input_column.success(f"This news is {label[prediction]}.")
             st.text(f"{int(result[prediction]*100)}% confidence")
-            input_column.progress(result[prediction])
+            input_column.progress(result[prediction])
+            
+        with st.spinner("Searching for references"):
+            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:
+                reference_column.write(f"""
+                <a href={data["url"][i]}><small>turnbackhoax.id</small></a>
+                <h2>{data["title"][i]}</h2>
+                """, unsafe_allow_html=True)
+                with reference_column.beta_expander("read content"):
+                    st.write(data["text"][i])