import bentoml
import yake


class KeywordExtractor(bentoml.Runnable):
    SUPPORTED_RESOURCES = ("cpu",)
    SUPPORTS_CPU_MULTI_THREADING = False

    def __init__(self, **kwargs):
        self.keyword_extractor = yake.KeywordExtractor(**kwargs)

    @bentoml.Runnable.method(batchable=False)
    def extract_keywords(self, transcript, lang, key_count):
        self.keyword_extractor.lan = lang
        self.keyword_extractor.top = key_count
        keywords = self.keyword_extractor.extract_keywords(transcript)
        print("keywords successfully extracted")
        return keywords