<commit_before>from typing import Dict, Generator, Iterable, Optional, Sequence, Union

import spacy

from .base import Detector
from ..filth import NamedEntityFilth, Filth
from ..utils import CanonicalStringSet


class NamedEntityDetector(Detector):
    """Use spacy's named entity recognition to clean named entities.
     List specific entities to include passing ``named_entities``, e.g.
     (PERSON)
    """
    filth_cls = NamedEntityFilth
    name = 'named_entity'

    disallowed_nouns = CanonicalStringSet(["skype"])

    def __init__(self, named_entities: Iterable[str] = {'PERSON'},
                 model: str = "en_core_web_trf", **kwargs):
        # Spacy NER are all upper cased
        self.named_entities = {entity.upper() for entity in named_entities}
        if model not in spacy.info()['pipelines']:
            raise OSError("Can't find model '{}'. If it is a valid Spacy model, "
                          "download it (e.g. with the CLI command "
                          "`python -m spacy download {}`).".format(model, model))
        self.nlp = spacy.load(model)
        # Only enable necessary pipes
        self.nlp.select_pipes(enable=["transformer", "tagger", "parser", "ner"])
        super(NamedEntityDetector, self).__init__(**kwargs)

    def _iter_spacy_pipeline(self, doc_names: Sequence[Optional[str]], doc_list: Sequence[str]):
        for doc_name, doc in zip(doc_names, self.nlp.pipe(doc_list)):
            for ent in doc.ents:
                if ent.label_ in self.named_entities:
                    yield self.filth_cls(beg=ent.start_char,
                                         end=ent.end_char,
                                         text=ent.text,
                                         document_name=None or str(doc_name),  # None if no doc_name provided
                                         detector_name=self.name,
                                         label=ent.label_)

    def iter_filth_documents(self, documents: Union[Sequence[str], Dict[str, str]]) -> Generator[Filth, None, None]:
        if isinstance(documents, list):
            doc_names, doc_list = zip(*enumerate(documents))
        elif isinstance(documents, dict):
            doc_names, doc_list = zip(*documents.items())
        else:
            raise TypeError('documents must be one of a string, list of strings or dict of strings.')

        yield from self._iter_spacy_pipeline(doc_names, doc_list)

    def iter_filth(self, text: str, document_name: Optional[str] = None) -> Generator[Filth, None, None]:
        yield from self._iter_spacy_pipeline([document_name], [text])<commit_msg>NamedEntityDetector to return standard Filth when it is avaliable<commit_after>from typing import Dict, Generator, Iterable, Optional, Sequence, Union

import spacy

from .base import Detector
from ..filth import NamedEntityFilth, Filth, NameFilth, OrganizationFilth
from ..utils import CanonicalStringSet


class NamedEntityDetector(Detector):
    """Use spacy's named entity recognition to clean named entities.
     List specific entities to include passing ``named_entities``, e.g.
     (PERSON)
    """
    filth_cls_map = {
        'PERSON': NameFilth,
        'ORG': OrganizationFilth
    }
    name = 'named_entity'

    disallowed_nouns = CanonicalStringSet(["skype"])

    def __init__(self, named_entities: Iterable[str] = {'PERSON'},
                 model: str = "en_core_web_trf", **kwargs):
        # Spacy NER are all upper cased
        self.named_entities = {entity.upper() for entity in named_entities}
        if model not in spacy.info()['pipelines']:
            raise OSError("Can't find model '{}'. If it is a valid Spacy model, "
                          "download it (e.g. with the CLI command "
                          "`python -m spacy download {}`).".format(model, model))
        self.nlp = spacy.load(model)
        # Only enable necessary pipes
        self.nlp.select_pipes(enable=["transformer", "tagger", "parser", "ner"])
        super(NamedEntityDetector, self).__init__(**kwargs)

    def _iter_spacy_pipeline(self, doc_names: Sequence[Optional[str]], doc_list: Sequence[str]):
        for doc_name, doc in zip(doc_names, self.nlp.pipe(doc_list)):
            print(doc_name)
            for ent in doc.ents:
                if ent.label_ in self.named_entities:
                    # If there is no standard 'filth', returns a NamedEntity filth
                    filth_cls = self.filth_cls_map.get(ent.label_, NamedEntityFilth)
                    yield filth_cls(beg=ent.start_char,
                                    end=ent.end_char,
                                    text=ent.text,
                                    document_name=(str(doc_name) if doc_name else None),  # None if no doc_name provided
                                    detector_name=self.name,
                                    label=ent.label_)

    def iter_filth_documents(self, documents: Union[Sequence[str], Dict[str, str]]) -> Generator[Filth, None, None]:
        if isinstance(documents, list):
            doc_names, doc_list = zip(*enumerate(documents))
        elif isinstance(documents, dict):
            doc_names, doc_list = zip(*documents.items())
        else:
            raise TypeError('documents must be one of a string, list of strings or dict of strings.')

        yield from self._iter_spacy_pipeline(doc_names, doc_list)

    def iter_filth(self, text: str, document_name: Optional[str] = None) -> Generator[Filth, None, None]:
        yield from self._iter_spacy_pipeline([document_name], [text])
{'fraction_non_alphanumeric': 0.10066391530593935, 'fraction_numerical': 0.0, 'mean_word_length': 2.3558097531607465, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 19, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1866440', 'n_tokens_mistral': 1598, 'n_tokens_neox': 1556, 'n_words': 430}