extracting keywords from texts

.gitignore

@@ -0,0 +1,2 @@
+/__pycache__*
+recursive-exclude * *.py[co]

app.py

@@ -1,14 +1,24 @@
 import streamlit as st
-from .generation import Insights
+from generation import Insights
+
 
 import time
 
 st.title("Drop the first document")
-file1 = st.file_uploader("Upload a file", type=["md", "txt"])
+file1 = st.file_uploader("Upload a file", type=["md", "txt"], key="first")
 st.title("Drop the second document")
-file2 = st.file_uploader("Upload a file", type=["md", "txt"])
+file2 = st.file_uploader("Upload a file", type=["md", "txt"], key="second")
 
-st.title("Contents of the first file")
-st.write(file1.read())
-st.title("Contents of the second file")
-st.write(file2.read())
+if file1 is not None and file2 is not None:
+    st.title("Contents of the first file")
+    st.title("Contents of the second file")
+    
+    st.title("Generating insights")
+    with st.spinner('Generating insights...'):
+        insight1 = Insights(file1.read().decode("utf-8"))
+        insight2 = Insights(file2.read().decode("utf-8"))
+        st.write(insight1.text)
+        st.write(insight2.text)
+        st.write(insight1.generate_topics())
+        st.write(insight2.generate_topics())
+        st.success('Done!')

generation.py

@@ -1,8 +1,25 @@
 import openai 
+from utils import *
+import mdforest 
 
 class Insights:
     
-    def __init__(self) -> None:
-        pass
+    EMBEDDING_MAX_TOKENS = 8191
+    
+    def __init__(self, text:str) -> None:
+        self.corpus = preprocess(mdforest.clean_markdown(text))
+        self.text = create_nest_sentences(self.corpus, self.EMBEDDING_MAX_TOKENS)
+        self.keywords = []
+        self.model = load_keyword_model()
+    
+    def generate_topics(self) -> list:
+        print("We are here for generating topics")
+        for sentence in self.text:
+            self.keywords = self.keywords  + generate_keywords(self.model, sentence)
+        return self.keywords
+    
+    
+    
+        
     
     

utils.py

@@ -1,14 +1,24 @@
 import streamlit as st
 from keybert import KeyBERT
+from nltk.corpus import stopwords
 from transformers import AutoTokenizer
 import re
 
+@st.cache_data
+def load_autotoken():
+  autotok = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')
+  return autotok
+
+@st.cache_data
+def load_keyword_model():
+  kw_model = KeyBERT()
+  return kw_model
 
-def create_nest_sentences(document:str, token_max_length = 1024):
+def create_nest_sentences(document:str, token_max_length = 8191):
   nested = []
   sent = []
   length = 0
-  tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')
+  tokenizer = load_autotoken()
 
   for sentence in re.split(r'(?<=[^A-Z].[.?]) +(?=[A-Z])', document.replace("\n", ' ')):
     tokens_in_sentence = tokenizer(str(sentence), truncation=False, padding=False)[0] # hugging face transformer tokenizer
@@ -25,19 +35,20 @@ def create_nest_sentences(document:str, token_max_length = 1024):
     nested.append(sent)
   return nested
 
-@st.cache_data
-def load_keyword_model():
-  kw_model = KeyBERT()
-  return kw_model
-
-
-def keyword_gen(kw_model, sequence:str):
-    keywords = kw_model.extract_keywords(
-        sequence, 
-        keyphrase_ngram_range=(1, 2),
-        stop_words='english', 
-        use_mmr=True, 
-        diversity=0.5,
-        top_n=10
-    )
-    return keywords
+def preprocess(text) -> str:
+    stop_words = set(stopwords.words("english"))
+    text = text.lower()
+    text = ''.join([c for c in text if c not in ('!', '.', ',', '?', ':', ';', '"', "'", '-', '(', ')')])
+    words = text.split()
+    words = [w for w in words if not w in stop_words]
+    return " ".join(words)
+
+def generate_keywords(kw_model, document: str) -> list:
+    atomic_extractions = kw_model.extract_keywords(document, keyphrase_ngram_range=(1, 1), stop_words=None, use_maxsum=True, nr_candidates=20, top_n=10)
+    complex_extractions = kw_model.extract_keywords(document, keyphrase_ngram_range=(1, 2), stop_words=None, use_maxsum=True, nr_candidates=20, top_n=10)
+    final_topics = []
+    for extraction in atomic_extractions:
+        final_topics.append(extraction[0])
+    for extraction in complex_extractions:
+        final_topics.append(extraction[0])
+    return final_topics