Adding app.py and pipeline.py, changed code structure

.gitignore

@@ -1,2 +1,4 @@
 venv
 test_file.txt
+scrap.py
+

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+import torch
+import os
+
+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):
+    inp = {"text": text, "title": title}
+    keyword_ls = kw_pipeline(inputs=inp, min_freq=min_freq, 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
+
+
+if gr.NO_RELOAD:
+    print("Loading PhoBERT model")
+    phobert = torch.load(f'{DIR_PATH}/pretrained-models/phobert.pt')
+    phobert.eval()
+
+    print("Loading NER model")
+    ner_model = torch.load(f'{DIR_PATH}/pretrained-models/ner-vietnamese-electra-base.pt')
+    ner_model.eval()
+    kw_pipeline = KeywordExtractorPipeline(phobert, ner_model)
+
+if __name__ == "__main__":
+    demo = gr.Interface(fn=extract_keyword,
+                        inputs=[
+                            gr.Text(
+                                label="Title",
+                                lines=1,
+                                value="Enter title here",
+                            ),
+                            gr.Textbox(
+                                label="Text",
+                                lines=5,
+                                value="Enter text here",
+                            ),
+                            gr.Number(
+                                label="Top N keywords",
+                                info="Number of keywords retrieved",
+                                value=10
+                            ),
+                            gr.Number(
+                                label="Ngram low range",
+                                value=1
+                            ),
+                            gr.Number(
+                                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",
+                        )
+                        )
+
+    demo.launch(share=True)  # Share your demo with just 1 extra parameter 🚀

keybertvi_model.py

@@ -0,0 +1,72 @@
+# import py_vncorenlp
+# from transformers import AutoTokenizer, pipeline
+# import torch
+# import os
+# from model.keyword_extraction_utils import extract_keywords
+#
+#
+# class KeyBERTVi:
+#
+#     def __init__(self, stopwords_file_path=None):
+#         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 model")
+#         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()
+#
+#         print("Loading NER model")
+#         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)
+#
+#         if stopwords_file_path is None:
+#             stopwords_file_path = f'{dir_path}/vietnamese-stopwords-dash.txt'
+#         with open(stopwords_file_path) as f:
+#             self.stopwords = [w.strip() for w in f.readlines()]
+#
+#     def extract_keywords(self, title, text, ngram_range=(1, 3), top_n=5, use_kmeans=False, use_mmr=False, min_freq=1):
+#         keyword_ls = extract_keywords(text, title,
+#                                       self.ner_pipeline,
+#                                       self.annotator,
+#                                       self.phobert_tokenizer,
+#                                       self.phobert,
+#                                       self.stopwords,
+#                                       ngram_n=ngram_range,
+#                                       top_n=top_n,
+#                                       use_kmeans=use_kmeans,
+#                                       use_mmr=use_mmr,
+#                                       min_freq=min_freq)
+#         return keyword_ls
+#
+#     def highlight(self, text, keywords):
+#         kw_ls = [' '.join(kw.split('_')) for kw, score in keywords]
+#         for key in kw_ls:
+#             text = text.replace(f" {key}", f" <mark>{key}</mark>")
+#         return text
+#
+#
+# dir_path = os.path.dirname(os.path.realpath(__file__))
+# if __name__ == "__main__":
+#     # args
+#     # print(dir_path)
+#
+#     stopwords_file_path = f'{dir_path}/vietnamese-stopwords-dash.txt'
+#
+#     # text_file_path = sys.argv[1]
+#     # with open(f'{dir_path}/{text_file_path}', 'r') as f:
+#     #     text = ' '.join([ln.strip() for ln in f.readlines()])
+#         # print(text)
+#
+#     # kw_model = KeyBERTVi()
+#     # model_name_on_hub = "KeyBERTVi"
+#     # kw_model.save_pretrained(model_name_on_hub)
+#     # kw_model.phobert_tokenizer.save_pretrained(model_name_on_hub)
+#
+#     # title = None
+#     # keyword_ls = kw_model.extract_keywords(title, text, ngram_range=(1, 3), top_n=5)
+#     # print(keyword_ls)

model.py

@@ -1,58 +0,0 @@
-import py_vncorenlp
-from transformers import AutoTokenizer, pipeline
-import torch
-import os
-from keyword_extraction import extract_keywords
-import sys
-
-
-class KeyBERTVi:
-
-    def __init__(self, stopwords_file_path):
-        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 model")
-        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()
-
-        print("Loading NER model")
-        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)
-
-        with open(stopwords_file_path) as f:
-            self.stopwords = [w.strip() for w in f.readlines()]
-
-    def extract_keywords(self, title, text, ngram_range=(1, 3), top_n=5):
-        keyword_ls = extract_keywords(text, title,
-                                      self.ner_pipeline,
-                                      self.annotator,
-                                      self.phobert_tokenizer,
-                                      self.phobert,
-                                      self.stopwords,
-                                      ngram_n=ngram_range,
-                                      top_n=top_n)
-        return keyword_ls
-
-
-dir_path = os.path.dirname(os.path.realpath(__file__))
-if __name__ == "__main__":
-    # args
-    # print(dir_path)
-
-    stopwords_file_path = f'{dir_path}/vietnamese-stopwords-dash.txt'
-
-    text_file_path = sys.argv[1]
-    with open(f'{dir_path}/{text_file_path}', 'r') as f:
-        text = ' '.join([ln.strip() for ln in f.readlines()])
-        # print(text)
-
-    kw_model = KeyBERTVi(stopwords_file_path)
-    title = None
-    keyword_ls = kw_model.extract_keywords(title, text, ngram_range=(1, 3), top_n=5)
-    print(keyword_ls)

keyword_extraction.py → model/keyword_extraction_utils.py

@@ -2,13 +2,11 @@ from string import punctuation
 import numpy as np
 import torch
 from sklearn.cluster import KMeans
-from named_entities import get_named_entities
+from model.named_entities import get_named_entities
 
 punctuation = [c for c in punctuation if c != "_"]
 punctuation += ["“", "–", ",", "…", "”", "–"]
 
-# with open('/Users/thaopham/bao-lao-dong-data/utils/vietnamese-stopwords-dash.txt') as f:
-#     stopwords = [w.strip() for w in f.readlines()]
 
 ethnicity_dict_map = {"H'Mông": "HMông",
                       "H'mông": "HMông",
@@ -69,8 +67,8 @@ def compute_ngram_list(segmentised_doc, ngram_n, stopwords_ls, subsentences=True
     ngram_list = []
     for sentence in output_sub_sentences:
         sent = sentence.split()
-        for i in range(len(sent) - (ngram_n) + 1):
-            ngram = ' '.join(sent[i:i + (ngram_n)])
+        for i in range(len(sent) - ngram_n + 1):
+            ngram = ' '.join(sent[i:i + ngram_n])
             if ngram not in ngram_list and not check_for_stopwords(ngram, stopwords_ls):
                 ngram_list.append(ngram)
 
@@ -116,15 +114,11 @@ def get_segmentised_doc(nlp, rdrsegmenter, title, doc):
             title = title.replace(i, j)
         doc = doc.replace(i, j)
 
-    print(title)
-    # print(doc)
     segmentised_doc = rdrsegmenter.word_segment(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))
-    print(sorted(ne_ls))
 
     segmentised_doc_ne = []
     for sent in segmentised_doc:
@@ -150,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 = {}
 
@@ -168,20 +155,7 @@ def compute_ngram_similarity(ngram_list, ngram_embeddings, doc_embedding):
     return ngram_similarity_dict
 
 
-# def compute_ngram_similarity(ngram_list, ngram_embeddings, doc_embedding):
-#     ngram_similarity_dict = {}
-#
-#     for ngram in ngram_list:
-#         score = [cosine_similarity(ngram_embeddings[ngram], doc_embedding.T).flatten()[0]]
-#         for w in ngram.split():
-#             similarity_score = cosine_similarity(ngram_embeddings[w], doc_embedding.T).flatten()[0]
-#             score.append(similarity_score)
-#         ngram_similarity_dict[ngram] = np.mean(score)
-#
-#     return ngram_similarity_dict
-
-
-def diversify_result(ngram_result, ngram_embeddings, top_n=5):
+def diversify_result_kmeans(ngram_result, ngram_embeddings, top_n=5):
     best_ngrams = sorted(ngram_result, key=ngram_result.get, reverse=True)[:top_n * 4]
     best_ngram_embeddings = np.array([ngram_embeddings[ngram] for ngram in best_ngrams]).squeeze()
     vote = {}
@@ -208,36 +182,8 @@ def diversify_result(ngram_result, ngram_embeddings, top_n=5):
     return diversify_result_ls[:top_n]
 
 
-def mmr(ngram_result, ngram_embeddings, lambda_=0.7, top_n=5):
-    ngram_result = {key: ngram_result[key] for key in
-                    sorted(ngram_result, key=ngram_result.get, reverse=True)[:top_n * 4]}
-
-    mmr_result = {}
-    for ngram1 in ngram_result:
-        similary_score_to_doc = ngram_result[ngram1]
-        max_sim = -1
-        # most_similar_ngram = None
-        for ngram2 in ngram_result:
-            if ngram2 != ngram1:
-                similarity_score_to_ngram = \
-                    cosine_similarity(ngram_embeddings[ngram1], ngram_embeddings[ngram2].T).flatten()[0]
-
-                if ngram2.lower() == ngram1.lower():
-                    similarity_score_to_ngram = 1
-
-                if similarity_score_to_ngram > max_sim:
-                    max_sim = similarity_score_to_ngram
-                    # most_similar_ngram = ngram2
-        # print(ngram1, most_similar_ngram, max_sim)
-        mmr_result[ngram1] = lambda_ * similary_score_to_doc - (1 - lambda_) * max_sim
-
-    mmr_result_ls = [(key, mmr_result[key]) for key in mmr_result]
-    mmr_result_ls = sorted(mmr_result_ls, key=lambda x: x[1], reverse=True)
-    return mmr_result_ls[:top_n]
-
-
 def remove_duplicates(ngram_result):
-    to_remove = []
+    to_remove = set()
     for ngram in ngram_result:
 
         for ngram2 in ngram_result:
@@ -245,7 +191,7 @@ def remove_duplicates(ngram_result):
                 new_score = np.mean([ngram_result[ngram], ngram_result[ngram2]])
 
                 ngram_result[ngram] = new_score
-                to_remove.append(ngram2)
+                to_remove.add(ngram2)
 
     for ngram in to_remove:
         ngram_result.pop(ngram)
@@ -257,9 +203,9 @@ 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', 'A']
-    keep_tags = ['N', 'Np', 'V']
+    keep_tags = ['N', 'Np', 'V', 'Nc']
     for key in annotated.keys():
+        # print(key,annotated[key])
         sent = ' '.join([dict_['wordForm'] for dict_ in annotated[key] if dict_['posTag'] in keep_tags])
         filtered_sentences.append(sent)
     return filtered_sentences
@@ -267,11 +213,9 @@ 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)
 
@@ -280,71 +224,37 @@ def get_candidate_ngrams(segmentised_doc, filtered_segmentised_doc, ngram_n, sto
     return candidate_ngram
 
 
-def get_ngram_frequencies(doc_segmentised, ngram_list):
+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()
+            # print(sent)
             for i in range(len(sent) - ngram_n + 1):
                 pair = ' '.join(sent[i:i + ngram_n])
+                # print(pair, ngram)
                 if pair == ngram:
                     count += 1
-        if count > 1:
+            # print(ngram, count)
+        if count >= min_freq:
             ngram_dict_freq[ngram] = count
 
     return ngram_dict_freq
 
 
-def extract_keywords(text, title, nlp, annotator, tokenizer, phobert, stopwords_ls, ngram_n=(2, 2), top_n=5,
-                     use_kmeans=False, use_mmr=False):
-    # ROOT_DIR = '/Users/thaopham/bao-lao-dong-data'
-    # with open(f'{ROOT_DIR}/utils/vietnamese-stopwords-dash.txt') as f:
-    #     stopwords = [w.strip() for w in f.readlines()]
-
-    ngram_low, ngram_high = ngram_n
-
-    ne_ls, doc_segmentised = get_segmentised_doc(nlp, annotator, title, text)
-    filtered_doc_segmentised = compute_filtered_text(annotator, title, text)
-    # print(doc_segmentised)
-
-    print("Generate ngram list")
-    ngram_list = set()
-    for n in range(ngram_low, ngram_high + 1):
-        # ngram_list += compute_ngram_list(segmentised_doc=doc_segmentised, ngram_n=n, stopwords_ls=stopwords)
-        # ngram_list.update(compute_ngram_list(segmentised_doc=doc_segmentised, ngram_n=n, stopwords_ls=stopwords))
-        ngram_list.update(get_candidate_ngrams(doc_segmentised, filtered_doc_segmentised, n, stopwords_ls))
-    ngram_list.update([annotator.word_segment(ne)[0] for ne in ne_ls])
-
-    print(len(ngram_list))
-    ngram_list = get_ngram_frequencies(doc_segmentised, ngram_list)
-    print(len(ngram_list))
-    # print(sorted(ngram_list))
-
-    print("Generate ngram embeddings")
-    ngram_embeddings = compute_ngram_embeddings(tokenizer, phobert, ngram_list)
-
-    print("Generate document embeddings")
-    doc_embedding = get_doc_embeddings(filtered_doc_segmentised, tokenizer, phobert, stopwords_ls)
-
-    ngram_result = compute_ngram_similarity(ngram_list, ngram_embeddings, doc_embedding)
-    ngram_result = remove_duplicates(ngram_result)
-    non_diversified = sorted([(ngram, ngram_result[ngram]) for ngram in ngram_result],
-                             key=lambda x: x[1], reverse=True)[:top_n]
-
-    # Diversify result
-    if use_kmeans:
-        diversified_kw_kmeans = diversify_result(ngram_result, ngram_embeddings, top_n=top_n)
-        return diversified_kw_kmeans
-
-    if use_mmr:
-        diversified_kw_mmr = mmr(ngram_result, ngram_embeddings, lambda_=0.85, top_n=top_n)
-        return diversified_kw_mmr
-    return non_diversified
-
-
-def highlight(text, keywords):
-    for key in keywords:
-        text = text.replace(f" {key}", f" <mark>{key}</mark>")
-    return text
+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)
+
+    # print("To removed")
+    # print(to_remove)
+    for kw in to_remove:
+        ngram_list.remove(kw)
+    return ngram_list

named_entities.py → model/named_entities.py

@@ -1,6 +1,4 @@
-# from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
 from underthesea import sent_tokenize
-import torch
 
 
 def substring(w, ls):

process_text.py → model/process_text.py

@@ -1,8 +1,7 @@
 from string import punctuation
-import re
 
 
-def process_article_content(text):
+def process_text_pipeline(text):
     full_text_processed = replace_all(text.strip())
 
     while '\n\n' in full_text_processed:
@@ -10,12 +9,6 @@ def process_article_content(text):
 
     full_text_processed = process_sticking_sentences(full_text_processed)
 
-    pattern = "Ảnh: [A-ZĐÀÁÂÃÈÉÊẾÌÍÒÓÔÕÙÚĂĐĨŨƠƯĂẠẢẤẦẨẪẬẮẰẲẴẶẸẺẼỀỂỄỆỈỊỌỎỐỒỔỖỘỚỜỞỠỢỤỦỨỪỬỮỰỲỴÝỶỸa-zđ][a-zàáâãèéêếìíòóôõùúăđĩũơưăạảấầẩẫậắằẳẵặẹẻẽềểễệỉịọỏốồổỗộớờởỡợụủứừửữựỳỵỷỹ]+ [A-ZĐÀÁÂÃÈÉÊẾÌÍÒÓÔÕÙÚĂĐĨŨƠƯĂẠẢẤẦẨẪẬẮẰẲẴẶẸẺẼỀỂỄỆỈỊỌỎỐỒỔỖỘỚỜỞỠỢỤỦỨỪỬỮỰỲỴÝỶỸa-zđ][a-zàáâãèéêếìíòóôõùúăđĩũơưăạảấầẩẫậắằẳẵặẹẻẽềểễệỉịọỏốồổỗộớờởỡợụủứừửữựỳỵỷỹ]+\.?"
-    full_text_processed = re.sub(pattern, '', full_text_processed)
-
-    pattern = "Ảnh: [A-ZĐÀÁÂÃÈÉÊẾÌÍÒÓÔÕÙÚĂĐĨŨƠƯĂẠẢẤẦẨẪẬẮẰẲẴẶẸẺẼỀỂỄỆỈỊỌỎỐỒỔỖỘỚỜỞỠỢỤỦỨỪỬỮỰỲỴÝỶỸ]+\.?"
-    full_text_processed = re.sub(pattern, '', full_text_processed)
-
     while '  ' in full_text_processed:
         full_text_processed = full_text_processed.replace('  ', ' ')
     return full_text_processed

pipeline.py

@@ -0,0 +1,134 @@
+import py_vncorenlp
+from transformers import AutoTokenizer, Pipeline, pipeline
+import os
+
+from model.keyword_extraction_utils import *
+from model.process_text import process_text_pipeline
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+
+
+class KeywordExtractorPipeline(Pipeline):
+    def __init__(self, model, ner_model, **kwargs):
+        super().__init__(model, **kwargs)
+        self.annotator = py_vncorenlp.VnCoreNLP(annotators=["wseg", "pos"],
+                                                save_dir=f'{dir_path}/pretrained-models/vncorenlp')
+
+        print("Loading PhoBERT tokenizer")
+        self.phobert_tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base-v2")
+        self.phobert = model
+
+        print("Loading NER tokenizer")
+        ner_tokenizer = AutoTokenizer.from_pretrained("NlpHUST/ner-vietnamese-electra-base")
+        self.ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer)
+
+        stopwords_file_path = f'{dir_path}/vietnamese-stopwords-dash.txt'
+        with open(stopwords_file_path) as f:
+            self.stopwords = [w.strip() for w in f.readlines()]
+
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {}
+        forward_kwargs = {}
+        postprocess_kwargs = {}
+
+        for possible_preprocess_kwarg in ["text", "title"]:
+            if possible_preprocess_kwarg in kwargs:
+                preprocess_kwargs[possible_preprocess_kwarg] = kwargs[possible_preprocess_kwarg]
+
+        for possible_forward_kwarg in ["ngram_n", "min_freq"]:
+            if possible_forward_kwarg in kwargs:
+                forward_kwargs[possible_forward_kwarg] = kwargs[possible_forward_kwarg]
+
+        for possible_postprocess_kwarg in ["top_n", "diversify_result"]:
+            if possible_postprocess_kwarg in kwargs:
+                postprocess_kwargs[possible_postprocess_kwarg] = kwargs[possible_postprocess_kwarg]
+
+        return preprocess_kwargs, forward_kwargs, postprocess_kwargs
+
+    def preprocess(self, inputs):
+        title = None
+        if inputs['title']:
+            title = process_text_pipeline(inputs['title'])
+        text = process_text_pipeline(inputs['text'])
+        return {"text": text, "title": title}
+
+    def _forward(self, model_inputs, ngram_n, min_freq):
+        text = model_inputs['text']
+        title = model_inputs['title']
+
+        # Getting segmentised document
+        ne_ls, doc_segmentised = get_segmentised_doc(self.ner_pipeline, self.annotator, title, text)
+        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("Final ngram list")
+        print(sorted(ngram_list))
+
+        ngram_embeddings = compute_ngram_embeddings(self.phobert_tokenizer, self.phobert, ngram_list)
+
+        return {"ngram_list": ngram_list, "ngram_embeddings": ngram_embeddings, "doc_embedding": doc_embedding}
+
+    def postprocess(self, model_outputs, top_n, diversify_result):
+        ngram_list = model_outputs['ngram_list']
+        ngram_embeddings = model_outputs['ngram_embeddings']
+        doc_embedding = model_outputs['doc_embedding']
+
+        ngram_result = self.extract_keywords(doc_embedding, ngram_list, ngram_embeddings)
+        non_diversified = sorted([(ngram, ngram_result[ngram]) for ngram in ngram_result],
+                                 key=lambda x: x[1], reverse=True)[:top_n]
+
+        if diversify_result:
+            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):
+        ngram_low, ngram_high = ngram_n
+
+        # Adding ngram
+        ngram_list = set()
+        for n in range(ngram_low, ngram_high + 1):
+            ngram_list.update(get_candidate_ngrams(doc_segmentised, filtered_doc_segmentised, n, self.stopwords))
+
+        # print(sorted(ngram_list))
+        # Adding named entities ngram list
+        ne_ls_segmented = [self.annotator.word_segment(ne)[0] for ne in ne_ls]
+        print("Named Entities list")
+        print(ne_ls_segmented)
+        ngram_list.update(ne_ls_segmented)
+
+        # print(sorted(ngram_list))
+        # Removing overlapping ngrams
+        ngram_list = remove_overlapping_ngrams(ngram_list)
+        # print("Removed overlapping ngrams")
+        # print(sorted(ngram_list))
+
+        # Limit ngrams by minimum frequency
+        if min_freq > 1:
+            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)
+        ngram_result = remove_duplicates(ngram_result)
+        return ngram_result
+
+
+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()])
+
+    inp = {"text": text, "title": None}
+    kws = kw_pipeline(inputs=inp, min_freq=1, ngram_n=(1, 3), top_n=5, diversify_result=False)
+    print(kws)

requirements.txt

@@ -0,0 +1,3 @@
+torch
+py_vncorenlp
+transformers

vietnamese-stopwords-dash.txt

@@ -868,7 +868,6 @@ ngày_xửa
 ngày_đến
 ngày_ấy
 ngôi
-ngôi_nhà
 ngôi_thứ
 ngõ_hầu
 ngăn_ngắt
@@ -1031,7 +1030,6 @@ năm
 năm_tháng
 nơi
 nơi_nơi
-nước
 nước_bài
 nước_cùng
 nước_lên
@@ -1987,4 +1985,14 @@ Nhờ
 Hiện
 Hiện_tại
 hiện
-cạnh
+cạnh
+Xung_quanh
+tóm_lại
+Cho
+Sau_này
+Vào
+ngày_nay
+chung_quanh
+Qua
+Thông_qua
+bao_gồm