| 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) |