Update app.py

app.py

@@ -3,6 +3,9 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipe
 import torch
 from pydantic import BaseModel
 from typing import Optional
+from sklearn.feature_extraction.text import CountVectorizer
+import yake
+
 
 app = FastAPI()
 
@@ -22,6 +25,15 @@ language_model = AutoModelForSequenceClassification.from_pretrained(model_name)
 language_tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
+language = "id"
+max_ngram_size = 3
+deduplication_threshold = 0.6
+deduplication_algo = 'seqm'
+windowSize = 3
+numOfKeywords = 20
+
+kw_extractor = yake.KeywordExtractor(lan=language, n=max_ngram_size, dedupLim=deduplication_threshold, dedupFunc=deduplication_algo, windowsSize=windowSize, top=numOfKeywords, features=None)
+
 
 
 
@@ -32,6 +44,37 @@ def greet_json():
     return {"Hello": "World!"}
 
 
+@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")
 async def language_detection(inp: InputText):