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 | |
| import uuid | |
| import json | |
| from transformers import AutoTokenizer, AutoModel | |
| from utils_demo import get_batch_text_representation | |
| base_dir = Path(__file__).parent | |
| class FHEAnonymizer: | |
| def __init__(self, punctuation_list=".,!?:;"): | |
| # Load tokenizer and model | |
| self.tokenizer = AutoTokenizer.from_pretrained("obi/deid_roberta_i2b2") | |
| self.embeddings_model = AutoModel.from_pretrained("obi/deid_roberta_i2b2") | |
| self.punctuation_list = punctuation_list | |
| with open(base_dir / "original_document_uuid_mapping.json", 'r') as file: | |
| self.uuid_map = json.load(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): | |
| # Pattern to identify words and non-words (including punctuation, spaces, etc.) | |
| token_pattern = r"(\b[\w\.\/\-@]+\b|[\s,.!?;:'\"-]+)" | |
| tokens = re.findall(token_pattern, text) | |
| identified_words_with_prob = [] | |
| processed_tokens = [] | |
| for token in tokens: | |
| # Directly append non-word tokens or whitespace to processed_tokens | |
| if not token.strip() or not re.match(r"\w+", token): | |
| processed_tokens.append(token) | |
| continue | |
| # Prediction for each word | |
| x = get_batch_text_representation([token], self.embeddings_model, self.tokenizer) | |
| prediction_proba = self.fhe_inference(x) | |
| probability = prediction_proba[0][1] | |
| if probability >= 0.5: | |
| identified_words_with_prob.append((token, probability)) | |
| # Use the existing UUID if available, otherwise generate a new one | |
| tmp_uuid = self.uuid_map.get(token, str(uuid.uuid4())[:8]) | |
| processed_tokens.append(tmp_uuid) | |
| self.uuid_map[token] = tmp_uuid | |
| else: | |
| processed_tokens.append(token) | |
| # Update the UUID map with query. | |
| with open(base_dir / "original_document_uuid_mapping.json", 'w') as file: | |
| json.dump(self.uuid_map, file) | |
| # Reconstruct the sentence | |
| reconstructed_sentence = ''.join(processed_tokens) | |
| return reconstructed_sentence, identified_words_with_prob | |