Add tool to get similar articles

app.py

@@ -1,5 +1,12 @@
+import torch
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import (
+    AutoModelForSeq2SeqLM, 
+    AutoTokenizer, 
+    AutoModelForTokenClassification,
+)
+import googlesearch
+
 
 def summarize(text):
     checkpoint = "sshleifer/distilbart-cnn-12-6"
@@ -28,9 +35,46 @@ def generate_question(text):
     return question
 
 
+def get_similar_articles(text):
+    tokenizer = AutoTokenizer.from_pretrained("yanekyuk/bert-keyword-extractor")
+    model = AutoModelForTokenClassification.from_pretrained("yanekyuk/bert-keyword-extractor")
+
+    inputs = tokenizer(text, truncation=True, return_tensors="pt")
+    outputs = model(**inputs)
+
+    keyword_tokens = []
+    current_keyword_tokens = []
+    for token, logits in zip(inputs.input_ids[0], outputs.logits[0]):
+        token_type = torch.argmax(logits).item()
+        if token_type > 0:
+            current_keyword_tokens.append(token.item())
+        elif len(current_keyword_tokens) > 0:
+            keyword_tokens.append(current_keyword_tokens)
+            current_keyword_tokens = []
+    
+    keywords = tokenizer.batch_decode(keyword_tokens)
+    keywords = list(set(keywords))
+
+    similar_websites = []
+    for keyword in keywords[:3]:
+        websites = googlesearch.search(
+            keyword,
+            tld="com",
+            lang="en",
+            num=3,
+            stop=3,
+            pause=0.5,
+        )
+        similar_websites += list(websites)
+    
+    return "\n".join(similar_websites)
+
+
 summarize_interface = gr.Interface(fn=summarize, inputs="text", outputs="text")
 question_interface = gr.Interface(fn=generate_question, inputs="text", outputs="text")
+similar_articles_interface = gr.Interface(fn=get_similar_articles, inputs="text", outputs="text")
 tabs = gr.TabbedInterface(
-    [summarize_interface, question_interface], ["Summarize an article", "Generate a question"]
+    [summarize_interface, question_interface, similar_articles_interface], 
+    ["Summarize an article", "Generate a question", "Get similar articles"],
 )
 tabs.launch()