File size: 1,639 Bytes
a2ee67b |
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 36 37 38 39 40 41 42 |
from transformers import Pipeline, AutoModelForTokenClassification, AutoTokenizer
class AnonymizationPipeline(Pipeline):
def __init__(self, model=None, tokenizer=None, **kwargs):
super().__init__(model=model, tokenizer=tokenizer, **kwargs)
if self.model is None:
self.model = AutoModelForTokenClassification.from_pretrained("JonathanEGP/Beto_Ner")
if self.tokenizer is None:
self.tokenizer = AutoTokenizer.from_pretrained("JonathanEGP/Beto_Ner")
self.ner_pipeline = Pipeline("ner", model=self.model, tokenizer=self.tokenizer)
def _sanitize_parameters(self, **kwargs):
return {}, {}, {} # No additional parameters needed for now
def preprocess(self, text):
return {"text": text}
def _forward(self, model_inputs):
text = model_inputs["text"]
entities = self.ner_pipeline(text)
return {"text": text, "entities": entities}
def postprocess(self, model_outputs):
text = model_outputs["text"]
entities = model_outputs["entities"]
# Ordenar las entidades de final a principio para no afectar los índices
entities.sort(key=lambda x: x['end'], reverse=True)
# Reemplazar las entidades con sus etiquetas
for entity in entities:
start = entity['start']
end = entity['end']
entity_type = entity['entity']
text = text[:start] + f"[{entity_type}]" + text[end:]
return {"anonymized_text": text}
def __call__(self, text, **kwargs):
return super().__call__(text, **kwargs) |