Spaces:
Sleeping
Sleeping
| import gensim | |
| import re | |
| from concrete.ml.deployment import FHEModelClient, FHEModelServer | |
| from pathlib import Path | |
| from concrete.ml.common.serialization.loaders import load | |
| base_dir = Path(__file__).parent | |
| class FHEAnonymizer: | |
| def __init__(self, punctuation_list=".,!?:;"): | |
| self.embeddings_model = gensim.models.FastText.load( | |
| str(base_dir / "embedded_model.model") | |
| ) | |
| self.punctuation_list = punctuation_list | |
| with open(base_dir / "cml_xgboost.model", "r") as model_file: | |
| self.fhe_ner_detection = load(file=model_file) | |
| path_to_model = (base_dir / "deployment").resolve() | |
| self.client = FHEModelClient(path_to_model) | |
| self.server = FHEModelServer(path_to_model) | |
| self.client.generate_private_and_evaluation_keys() | |
| self.evaluation_key = self.client.get_serialized_evaluation_keys() | |
| def fhe_inference(self, x): | |
| enc_x = self.client.quantize_encrypt_serialize(x) | |
| enc_y = self.server.run(enc_x, self.evaluation_key) | |
| y = self.client.deserialize_decrypt_dequantize(enc_y) | |
| return y | |
| def __call__(self, text: str): | |
| text = self.preprocess_sentences(text) | |
| identified_words_with_prob = [] # tuples of (word, probability) | |
| new_text = [] | |
| for word in text.split(): | |
| # Prediction for each word | |
| x = self.embeddings_model.wv[word][None] | |
| prediction_proba = self.fhe_ner_detection.predict_proba(x) | |
| # prediction = self.fhe_inference(x).argmax(1)[0] | |
| # print(word, prediction) | |
| probability = prediction_proba[0][1] | |
| prediction = probability >= 0.5 | |
| if prediction == 1: | |
| identified_words_with_prob.append((word, probability)) | |
| new_text.append("<REMOVED>") | |
| else: | |
| new_text.append(word) | |
| # Joining the modified text | |
| modified_text = " ".join(new_text) | |
| return modified_text, identified_words_with_prob | |
| def preprocess_sentences(self, sentence, verbose=False): | |
| """Preprocess the sentence.""" | |
| sentence = re.sub(r"\n+", " ", sentence) | |
| if verbose: | |
| print(sentence) | |
| sentence = re.sub(" +", " ", sentence) | |
| if verbose: | |
| print(sentence) | |
| sentence = re.sub(r"'s\b", " s", sentence) | |
| if verbose: | |
| print(sentence) | |
| sentence = re.sub(r"\s([,.!?;:])", r"\1", sentence) | |
| if verbose: | |
| print(sentence) | |
| pattern = r"(?<!\w)[{}]|[{}](?!\w)".format( | |
| re.escape(self.punctuation_list), re.escape(self.punctuation_list) | |
| ) | |
| sentence = re.sub(pattern, "", sentence) | |
| if verbose: | |
| print(sentence) | |
| sentence = re.sub(r"\s([,.!?;:])", r"\1", sentence) | |
| if verbose: | |
| print(sentence) | |
| return sentence | |