import spacy
import nltk

class NLPModel:
    def __init__(self):
        self.nlp = spacy.load("pt_core_news_md")
        nltk.download('punkt')
        
    def __call__(self, text: str):
        """Makes the model callable like model(text)."""
        return self.extract_entities(text)  # or another default method
    
    def extract_entities(self, text: str):
        if isinstance(text, list):  # If input is a list of sentences
            entities = []
            for sentence in text:
                doc = self.nlp(sentence)
                entities.extend([(ent.text.lower(), ent.label_) for ent in doc.ents])
            return entities
        else:  # If input is a single string
            doc = self.nlp(text)
            return [(ent.text.lower(), ent.label_) for ent in doc.ents]
    
    
    def tokenize_sentences(self, text: str):
        return nltk.sent_tokenize(text)