model/ner.py

from model.model import Model

class NER(Model):
    def __init__(self) -> None:
        self.model_name = None
        self.tasks = "ner"
        self.load_model(model_name="syafiqfaray/indobert-model-ner", tasks=self.tasks)
        
    def predict(self, sentences):
        outputs = super().predict(sentences, self.tasks)
        return {"result": outputs}
    
name_entity = NER()