Update app.py

app.py

@@ -35,7 +35,12 @@ numOfKeywords = 20
 kw_extractor = yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold, dedupFunc=deduplication_algo, windowsSize=windowSize, top=numOfKeywords, features=None)
 
 
-
+ner_model = "syafiqfaray/indobert-model-ner"
+ner = pipeline(
+    "ner", 
+    ner_model, 
+    aggregation_strategy="simple",
+)
 
 
 
@@ -45,35 +50,39 @@ def greet_json():
 
 
 @app.post("/key_phrase_extraction")
-async def key_phrase_extraction(inp:InputText):
-    def merge_keyphrases(keyphrases):
-        new_merged = keyphrases
-        while True:
-          merged = [new_merged[0]]
-          for i in range(1, len(keyphrases)):
-              keys = keyphrases[i]
-              keys_prev = keyphrases[i-1]
-              label = keys["label"]
-              score = keys["score"]
-              vectorizer = CountVectorizer(ngram_range=( 1,len(label.split(" ")) ), lowercase=False)
-              analyzer = vectorizer.build_analyzer()
-              for key in analyzer(label)[::-1]:
-                  key_prev = keys_prev["label"][::-1]
-                  if key == key_prev[:len(key)][::-1].strip():
-                      label = key_prev[len(key):][::-1].strip() + " " + label#.replace(key, "")
-                      score = max(keys_prev["score"],keys["score"])
-                      merged.pop()
-                      break
-              merged.append({"label":label.strip(), "score":score})
-          if new_merged == merged:
-            break
-          else:
-            new_merged = merged
-        return merged
-    
-    keywords = kw_extractor.extract_keywords(inp.text)
-    
-    return merge_keyphrases([{"label":key[0], "score":1-key[1]} for key in keywords if 1-key[1]>inp.threshold])
+async def key_phrase_extraction(inp: InputText):
+    return [{"label": x["word"], "score": x["score"]} for x in ner(inp.text) if x["score"] > inp.threshold]
+
+# @app.post("/key_phrase_extraction")
+# async def key_phrase_extraction(inp:InputText):
+#     def merge_keyphrases(keyphrases):
+#         new_merged = keyphrases
+#         while True:
+#           merged = [new_merged[0]]
+#           for i in range(1, len(keyphrases)):
+#               keys = keyphrases[i]
+#               keys_prev = keyphrases[i-1]
+#               label = keys["label"]
+#               score = keys["score"]
+#               vectorizer = CountVectorizer(ngram_range=( 1,len(label.split(" ")) ), lowercase=False)
+#               analyzer = vectorizer.build_analyzer()
+#               for key in analyzer(label)[::-1]:
+#                   key_prev = keys_prev["label"][::-1]
+#                   if key == key_prev[:len(key)][::-1].strip():
+#                       label = key_prev[len(key):][::-1].strip() + " " + label#.replace(key, "")
+#                       score = max(keys_prev["score"],keys["score"])
+#                       merged.pop()
+#                       break
+#               merged.append({"label":label.strip(), "score":score})
+#           if new_merged == merged:
+#             break
+#           else:
+#             new_merged = merged
+#         return merged
+    
+#     keywords = kw_extractor.extract_keywords(inp.text)
+    
+#     return merge_keyphrases([{"label":key[0], "score":1-key[1]} for key in keywords if 1-key[1]>inp.threshold])
 
 
 @app.post("/language_detection")