File size: 1,224 Bytes
a2682b3 e21244d a2682b3 e21244d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
import spacy
from typing import List, Union
import logging
logger = logging.getLogger(__name__)
class NLPModel:
def __init__(self):
try:
# Load spaCy model only
self.nlp = spacy.load("pt_core_news_md")
logger.info("spaCy model initialized successfully")
except Exception as e:
logger.error(f"Failed to initialize spaCy model: {str(e)}")
raise
def extract_entities(self, text: Union[str, List[str]]) -> List[tuple]:
"""Entity extraction using spaCy"""
try:
if isinstance(text, list):
text = " ".join(text)
doc = self.nlp(text)
return [(ent.text.lower(), ent.label_) for ent in doc.ents]
except Exception as e:
logger.error(f"Entity extraction failed: {str(e)}")
return []
def tokenize_sentences(self, text: str) -> List[str]:
"""Sentence tokenization using spaCy"""
try:
doc = self.nlp(text)
return [sent.text for sent in doc.sents]
except Exception as e:
logger.error(f"Sentence tokenization failed: {str(e)}")
return [text] # Fallback to returning whole text |