Remove minimum frequency and add highlighted output in app.py

app.py

@@ -2,22 +2,40 @@ import gradio as gr
 import torch
 import os
 
+from model.process_text import process_text_pipeline
 from pipeline import KeywordExtractorPipeline
 
 DIR_PATH = os.path.dirname(os.path.realpath(__file__))
 
 
-def extract_keyword(title, text, top_n, ngram_low_range, ngram_high_range, min_freq, diversify_result):
-    # title = None
-    # keyword_ls = kw_model.extract_keywords(title, text, ngram_range=(ngram_low_range, ngram_high_range), top_n=top_n,
-    #                                        min_freq=min_freq, use_kmeans=use_kmeans)
+def extract_keyword(title, text, top_n, ngram_low_range, ngram_high_range, diversify_result):
     inp = {"text": text, "title": title}
-    keyword_ls = kw_pipeline(inputs=inp, min_freq=min_freq, ngram_n=(ngram_low_range, ngram_high_range),
+    keyword_ls = kw_pipeline(inputs=inp, min_freq=1, ngram_n=(ngram_low_range, ngram_high_range),
                              top_n=top_n, diversify_result=diversify_result)
     result = ''
     for kw, score in keyword_ls:
         result += f'{kw}: {score}\n'
-    return result
+    highlight_dict = get_highlighted_text(text, keyword_ls)
+    return result, highlight_dict
+
+
+def get_highlighted_text(text, keyword_ls):
+    text = process_text_pipeline(text)
+    entities = []
+    for kw, score in keyword_ls:
+
+        kw = kw.replace('_', ' ')
+        kw_dict = {"entity": "KEY"}
+        start_idx = text.find(kw)
+
+        while start_idx != -1:
+            end_idx = start_idx + len(kw)
+            kw_dict["start"] = start_idx
+            kw_dict["end"] = end_idx
+            entities.append(kw_dict)
+            start_idx = text.find(kw, end_idx)
+
+    return {"text": text, "entities": entities}
 
 
 if gr.NO_RELOAD:
@@ -36,12 +54,14 @@ if __name__ == "__main__":
                             gr.Text(
                                 label="Title",
                                 lines=1,
-                                value="Enter title here",
+                                value="Truyền thuyết và hiện tại Thành Cổ Loa",
                             ),
                             gr.Textbox(
                                 label="Text",
                                 lines=5,
-                                value="Enter text here",
+                                value="""Nhắc đến Cổ Loa, người ta nghĩ ngay đến truyền thuyết về An Dương Vương được thần Kim Quy bày cho cách xây thành, về chiếc lẫy nỏ thần làm từ móng chân rùa thần và mối tình bi thương Mỵ Châu – Trọng Thủy. Đằng sau những câu chuyện thiên về tâm linh ấy, thế hệ con cháu còn khám phá được những giá trị khảo cổ to lớn của Cổ Loa.
+Khu di tích Cổ Loa cách trung – tâm Hà Nội 17km thuộc huyện Đông Anh, Hà Nội, có diện tích bảo tồn gần 500ha được coi là địa chỉ văn hóa đặc biệt của thủ đô và cả nước. Cổ Loa có hàng loạt di chỉ khảo cổ học đã được phát hiện, phản ánh quá trình phát triển liên tục của dân tộc ta từ sơ khai qua các thời kỳ đồ đồng, đồ đá và đồ sắt mà đỉnh cao là văn hóa Đông Sơn, vẫn được coi là nền văn minh sông Hồng thời kỳ tiền sử của dân tộc Việt Nam.
+Cổ Loa từng là kinh đô của nhà nước Âu Lạc thời kỳ An Dương Vương (thế kỷ III TCN) và của nước Đại Việt thời Ngô Quyền (thế kỷ X) mà thành Cổ Loa là một di tích minh chứng còn lại cho đến ngày nay. Thành Cổ Loa được các nhà khảo cổ học đánh giá là “tòa thành cổ nhất, quy mô lớn vào bậc nhất, cấu trúc cũng thuộc loại độc đáo nhất trong lịch sử xây dựng thành lũy của người Việt cổ”.""",
                             ),
                             gr.Number(
                                 label="Top N keywords",
@@ -56,18 +76,12 @@ if __name__ == "__main__":
                                 label="Ngram high range",
                                 value=3
                             ),
-                            gr.Number(
-                                label="Ngram minimum frequency",
-                                value=1
-                            ),
                             gr.Checkbox(
                                 label="Diversify result"
                             )
                         ],
-                        # inputs=["text", "textbox", "number", "number", "number", "number", "checkbox"],
-                        outputs=gr.Textbox(
-                            label="Keywords Extracted",
-                        )
+                        outputs=[gr.Textbox(label="Keywords Extracted"),
+                                 gr.HighlightedText(label="Highlight")]
                         )
 
     demo.launch(share=True)  # Share your demo with just 1 extra parameter 🚀

model/keyword_extraction_utils.py

@@ -118,7 +118,6 @@ def get_segmentised_doc(nlp, rdrsegmenter, title, doc):
 
     if title is not None:
         segmentised_doc = rdrsegmenter.word_segment(title) + rdrsegmenter.word_segment(doc)
-    print(segmentised_doc)
     ne_ls = set(get_named_entities(nlp, doc))
 
     segmentised_doc_ne = []
@@ -145,13 +144,6 @@ def compute_ngram_embeddings(tokenizer, phobert, ngram_list):
     return ngram_embeddings
 
 
-# def normalised_cosine_similarity(ngram_embedding, document_embedding):
-#     similarity_score = cosine_similarity(ngram_embedding, document_embedding)
-#     magnitude_ngram = np.linalg.norm(ngram_embedding)
-#     magnitude_doc = np.linalg.norm(document_embedding)
-#     return similarity_score / np.sqrt(magnitude_ngram * magnitude_doc)
-
-
 def compute_ngram_similarity(ngram_list, ngram_embeddings, doc_embedding):
     ngram_similarity_dict = {}
 
@@ -211,7 +203,7 @@ def compute_filtered_text(annotator, title, text):
     if title is not None:
         annotated = annotator.annotate_text(title + '. ' + text)
     filtered_sentences = []
-    keep_tags = ['N', 'Np', 'V']
+    keep_tags = ['N', 'Np', 'V', 'Nc']
     for key in annotated.keys():
         sent = ' '.join([dict_['wordForm'] for dict_ in annotated[key] if dict_['posTag'] in keep_tags])
         filtered_sentences.append(sent)
@@ -220,43 +212,22 @@ def compute_filtered_text(annotator, title, text):
 
 def get_candidate_ngrams(segmentised_doc, filtered_segmentised_doc, ngram_n, stopwords_ls):
     # get actual ngrams
-    # segmentised_doc = get_segmentised_doc(nlp, annotator, title, text)
     actual_ngram_list = compute_ngram_list(segmentised_doc, ngram_n, stopwords_ls, subsentences=True)
 
     # get filtered ngrams
-    # filtered_segmentised_doc = compute_filtered_text(annotator, title, text)
     filtered_ngram_list = compute_ngram_list(filtered_segmentised_doc, ngram_n, stopwords_ls,
                                              subsentences=False)
 
-    # get candiate ngrams
+    # get candidate ngrams
     candidate_ngram = [ngram for ngram in filtered_ngram_list if ngram in actual_ngram_list]
     return candidate_ngram
 
 
-def limit_minimum_frequency(doc_segmentised, ngram_list, min_freq=1):
-    ngram_dict_freq = {}
-    for ngram in ngram_list:
-        ngram_n = len(ngram.split())
-        count = 0
-        for sentence in doc_segmentised:
-            sent = sentence.split()
-            for i in range(len(sent) - ngram_n + 1):
-                pair = ' '.join(sent[i:i + ngram_n])
-                if pair == ngram:
-                    count += 1
-        if count >= min_freq:
-            ngram_dict_freq[ngram] = count
-
-    return ngram_dict_freq
-
-
 def remove_overlapping_ngrams(ngram_list):
     to_remove = set()
     for ngram1 in ngram_list:
         for ngram2 in ngram_list:
             if len(ngram1.split()) > len(ngram2.split()) and (ngram1.startswith(ngram2) or ngram1.endswith(ngram2)):
-                # print(ngram1, ngram2)
-                # print()
                 to_remove.add(ngram2)
 
     for kw in to_remove:

model/named_entities.py

@@ -32,8 +32,6 @@ def get_named_entities(nlp, doc):
         if len(sent_ner_result) > 0:
             ner_lists += get_ner_phrases(sent_ner_result)
 
-    # print(ner_lists)
-
     ner_list_non_dup = []
     for (entity, ner_type) in ner_lists:
         if entity not in ner_list_non_dup and ner_type.startswith('I'):

pipeline.py

@@ -13,19 +13,13 @@ class KeywordExtractorPipeline(Pipeline):
         super().__init__(model, **kwargs)
         self.annotator = py_vncorenlp.VnCoreNLP(annotators=["wseg", "pos"],
                                                 save_dir=f'{dir_path}/pretrained-models/vncorenlp')
-        # model = py_vncorenlp.VnCoreNLP(save_dir='/absolute/path/to/vncorenlp')
+
         print("Loading PhoBERT tokenizer")
         self.phobert_tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base-v2")
-
-        # use absolute path because torch is cached
-        # self.phobert = torch.load(f'{dir_path}/pretrained-models/phobert.pt')
-        # self.phobert.eval()
         self.phobert = model
 
         print("Loading NER tokenizer")
         ner_tokenizer = AutoTokenizer.from_pretrained("NlpHUST/ner-vietnamese-electra-base")
-        # ner_model = torch.load(f'{dir_path}/pretrained-models/ner-vietnamese-electra-base.pt')
-        # ner_model.eval()
         self.ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer)
 
         stopwords_file_path = f'{dir_path}/vietnamese-stopwords-dash.txt'
@@ -41,7 +35,7 @@ class KeywordExtractorPipeline(Pipeline):
             if possible_preprocess_kwarg in kwargs:
                 preprocess_kwargs[possible_preprocess_kwarg] = kwargs[possible_preprocess_kwarg]
 
-        for possible_forward_kwarg in ["ngram_n", "min_freq"]:
+        for possible_forward_kwarg in ["ngram_n"]:
             if possible_forward_kwarg in kwargs:
                 forward_kwargs[possible_forward_kwarg] = kwargs[possible_forward_kwarg]
 
@@ -49,8 +43,6 @@ class KeywordExtractorPipeline(Pipeline):
             if possible_postprocess_kwarg in kwargs:
                 postprocess_kwargs[possible_postprocess_kwarg] = kwargs[possible_postprocess_kwarg]
 
-        # print(forward_kwargs)
-        # print(postprocess_kwargs)
         return preprocess_kwargs, forward_kwargs, postprocess_kwargs
 
     def preprocess(self, inputs):
@@ -58,24 +50,22 @@ class KeywordExtractorPipeline(Pipeline):
         if inputs['title']:
             title = process_text_pipeline(inputs['title'])
         text = process_text_pipeline(inputs['text'])
-        return {"text": text, "title": title}#, "ngram_n": inputs["ngram_n"], "min_freq":inputs["min_freq"]}
+        return {"text": text, "title": title}
 
-    def _forward(self, model_inputs, ngram_n, min_freq):
+    def _forward(self, model_inputs, ngram_n):
         text = model_inputs['text']
         title = model_inputs['title']
-        # ngram_n = model_inputs['ngram_n']
-        # min_freq = model_inputs['min_freq']
 
         # Getting segmentised document
         ne_ls, doc_segmentised = get_segmentised_doc(self.ner_pipeline, self.annotator, title, text)
-        print(ne_ls)
         filtered_doc_segmentised = compute_filtered_text(self.annotator, title, text)
 
         doc_embedding = get_doc_embeddings(filtered_doc_segmentised, self.phobert_tokenizer, self.phobert,
                                            self.stopwords)
 
-        ngram_list = self.generate_ngram_list(doc_segmentised, filtered_doc_segmentised, ne_ls, ngram_n, min_freq)
-        print(sorted(ngram_list))
+        ngram_list = self.generate_ngram_list(doc_segmentised, filtered_doc_segmentised, ne_ls, ngram_n)
+        # print("Final ngram list")
+        # print(sorted(ngram_list))
 
         ngram_embeddings = compute_ngram_embeddings(self.phobert_tokenizer, self.phobert, ngram_list)
 
@@ -94,7 +84,7 @@ class KeywordExtractorPipeline(Pipeline):
             return diversify_result_kmeans(ngram_result, ngram_embeddings, top_n=top_n)
         return non_diversified
 
-    def generate_ngram_list(self, doc_segmentised, filtered_doc_segmentised, ne_ls, ngram_n, min_freq):
+    def generate_ngram_list(self, doc_segmentised, filtered_doc_segmentised, ne_ls, ngram_n):
         ngram_low, ngram_high = ngram_n
 
         # Adding ngram
@@ -103,14 +93,13 @@ class KeywordExtractorPipeline(Pipeline):
             ngram_list.update(get_candidate_ngrams(doc_segmentised, filtered_doc_segmentised, n, self.stopwords))
 
         # Adding named entities ngram list
-        ngram_list.update([self.annotator.word_segment(ne)[0] for ne in ne_ls])
+        ne_ls_segmented = [self.annotator.word_segment(ne)[0] for ne in ne_ls]
+        ngram_list.update(ne_ls_segmented)
 
         # Removing overlapping ngrams
         ngram_list = remove_overlapping_ngrams(ngram_list)
 
-        # Limit ngrams by minimum frequency
-        ngram_list = limit_minimum_frequency(doc_segmentised, ngram_list, min_freq=min_freq)
-        return ngram_list.keys()
+        return ngram_list
 
     def extract_keywords(self, doc_embedding, ngram_list, ngram_embeddings):
         ngram_result = compute_ngram_similarity(ngram_list, ngram_embeddings, doc_embedding)
@@ -119,17 +108,19 @@ class KeywordExtractorPipeline(Pipeline):
 
 
 if __name__ == "__main__":
-
     phobert = torch.load(f'{dir_path}/pretrained-models/phobert.pt')
     phobert.eval()
     ner_model = torch.load(f'{dir_path}/pretrained-models/ner-vietnamese-electra-base.pt')
     ner_model.eval()
     kw_pipeline = KeywordExtractorPipeline(phobert, ner_model)
 
-    text_file_path = f'{dir_path}/test_file.txt'
-    with open(text_file_path, 'r') as f:
-        text = ' '.join([ln.strip() for ln in f.readlines()])
+    title = "Truyền thuyết và hiện tại Thành Cổ Loa"
+    text = """
+            Nhắc đến Cổ Loa, người ta nghĩ ngay đến truyền thuyết về An Dương Vương được thần Kim Quy bày cho cách xây thành, về chiếc lẫy nỏ thần làm từ móng chân rùa thần và mối tình bi thương Mỵ Châu – Trọng Thủy. Đằng sau những câu chuyện thiên về tâm linh ấy, thế hệ con cháu còn khám phá được những giá trị khảo cổ to lớn của Cổ Loa.
+            Khu di tích Cổ Loa cách trung – tâm Hà Nội 17km thuộc huyện Đông Anh, Hà Nội, có diện tích bảo tồn gần 500ha được coi là địa chỉ văn hóa đặc biệt của thủ đô và cả nước. Cổ Loa có hàng loạt di chỉ khảo cổ học đã được phát hiện, phản ánh quá trình phát triển liên tục của dân tộc ta từ sơ khai qua các thời kỳ đồ đồng, đồ đá và đồ sắt mà đỉnh cao là văn hóa Đông Sơn, vẫn được coi là nền văn minh sông Hồng thời kỳ tiền sử của dân tộc Việt Nam.
+            Cổ Loa từng là kinh đô của nhà nước Âu Lạc thời kỳ An Dương Vương (thế kỷ III TCN) và của nước Đại Việt thời Ngô Quyền (thế kỷ X) mà thành Cổ Loa là một di tích minh chứng còn lại cho đến ngày nay. Thành Cổ Loa được các nhà khảo cổ học đánh giá là “tòa thành cổ nhất, quy mô lớn vào bậc nhất, cấu trúc cũng thuộc loại độc đáo nhất trong lịch sử xây dựng thành lũy của người Việt cổ”.
+        """
 
-    inp = {"text": text,"title": None}
-    kws = kw_pipeline(inputs=inp, min_freq=1, ngram_n=(1,3), top_n=5, diversify_result=False)
+    inp = {"text": text, "title": title}
+    kws = kw_pipeline(inputs=inp, ngram_n=(1, 3), top_n=5, diversify_result=False)
     print(kws)

requirements.txt

@@ -4,4 +4,4 @@ gradio
 scikit-learn
 numpy
 underthesea
-py_vncorenlp
+py_vncorenlp