bradarrML/encrypted_sentiment_analysis
0
1"""A gradio app. that runs locally (analytics=False and share=False) about sentiment analysis on tweets."""2 3import gradio as gr4from requests import head5from transformer_vectorizer import TransformerVectorizer6from concrete.ml.deployment import FHEModelClient7import numpy8import os9from pathlib import Path10import requests11import json12import base6413import subprocess14import shutil15import time16 17subprocess.Popen(["uvicorn", "server:app"])18 19# Wait 5 sec for the server to start20time.sleep(5)21 22# Encrypted data limit for the browser to display23# (encrypted data is too large to display in the browser)24ENCRYPTED_DATA_BROWSER_LIMIT = 50025N_USER_KEY_STORED = 2026 27print("Loading the transformer model...")28 29# Initialize the transformer vectorizer30transformer_vectorizer = TransformerVectorizer()31 32def clean_tmp_directory():33 # Allow 20 user keys to be stored.34 # Once that limitation is reached, deleted the oldest.35 path_sub_directories = sorted([f for f in Path(".fhe_keys/").iterdir() if f.is_dir()], key=os.path.getmtime)36 37 user_ids = []38 if len(path_sub_directories) > N_USER_KEY_STORED:39 n_files_to_delete = len(path_sub_directories) - N_USER_KEY_STORED40 for p in path_sub_directories[:n_files_to_delete]:41 user_ids.append(p.name)42 shutil.rmtree(p)43 44 list_files_tmp = Path("tmp/").iterdir()45 # Delete all files related to user_id46 for file in list_files_tmp:47 for user_id in user_ids:48 if file.name.endswith(f"{user_id}.npy"):49 file.unlink()50 51 52def keygen():53 # Clean tmp directory if needed54 clean_tmp_directory()55 56 print("Initializing FHEModelClient...")57 58 # Let's create a user_id59 user_id = numpy.random.randint(0, 2**32)60 fhe_api = FHEModelClient("sentiment_fhe_model/deployment", f".fhe_keys/{user_id}")61 fhe_api.load()62 63 64 # Generate a fresh key65 fhe_api.generate_private_and_evaluation_keys(force=True)66 evaluation_key = fhe_api.get_serialized_evaluation_keys()67 size_evaluation_key = len(evaluation_key)68 69 # Save evaluation_key in a file, since too large to pass through regular Gradio70 # buttons, https://github.com/gradio-app/gradio/issues/187771 numpy.save(f"tmp/tmp_evaluation_key_{user_id}.npy", evaluation_key)72 73 return [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], size_evaluation_key, user_id]74 75 76def encode_quantize_encrypt(text, user_id):77 if not user_id:78 raise gr.Error("You need to generate FHE keys first.")79 80 fhe_api = FHEModelClient("sentiment_fhe_model/deployment", f".fhe_keys/{user_id}")81 fhe_api.load()82 encodings = transformer_vectorizer.transform([text])83 quantized_encodings = fhe_api.model.quantize_input(encodings).astype(numpy.uint8)84 encrypted_quantized_encoding = fhe_api.quantize_encrypt_serialize(encodings)85 86 # Save encrypted_quantized_encoding in a file, since too large to pass through regular Gradio87 # buttons, https://github.com/gradio-app/gradio/issues/187788 numpy.save(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy", encrypted_quantized_encoding)89 90 # Compute size91 text_size = len(text.encode())92 encodings_size = len(encodings.tobytes())93 quantized_encoding_size = len(quantized_encodings.tobytes())94 encrypted_quantized_encoding_size = len(encrypted_quantized_encoding)95 encrypted_quantized_encoding_shorten = list(encrypted_quantized_encoding)[:ENCRYPTED_DATA_BROWSER_LIMIT]96 encrypted_quantized_encoding_shorten_hex = ''.join(f'{i:02x}' for i in encrypted_quantized_encoding_shorten)97 return (98 encodings[0],99 quantized_encodings[0],100 encrypted_quantized_encoding_shorten_hex,101 text_size,102 encodings_size,103 quantized_encoding_size,104 encrypted_quantized_encoding_size,105 )106 107 108def run_fhe(user_id):109 encoded_data_path = Path(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy")110 if not user_id:111 raise gr.Error("You need to generate FHE keys first.")112 if not encoded_data_path.is_file():113 raise gr.Error("No encrypted data was found. Encrypt the data before trying to predict.")114 115 # Read encrypted_quantized_encoding from the file116 encrypted_quantized_encoding = numpy.load(encoded_data_path)117 118 # Read evaluation_key from the file119 evaluation_key = numpy.load(f"tmp/tmp_evaluation_key_{user_id}.npy")120 121 # Use base64 to encode the encodings and evaluation key122 encrypted_quantized_encoding = base64.b64encode(encrypted_quantized_encoding).decode()123 encoded_evaluation_key = base64.b64encode(evaluation_key).decode()124 125 query = {}126 query["evaluation_key"] = encoded_evaluation_key127 query["encrypted_encoding"] = encrypted_quantized_encoding128 headers = {"Content-type": "application/json"}129 response = requests.post(130 "http://localhost:8000/predict_sentiment", data=json.dumps(query), headers=headers131 )132 encrypted_prediction = base64.b64decode(response.json()["encrypted_prediction"])133 134 # Save encrypted_prediction in a file, since too large to pass through regular Gradio135 # buttons, https://github.com/gradio-app/gradio/issues/1877136 numpy.save(f"tmp/tmp_encrypted_prediction_{user_id}.npy", encrypted_prediction)137 encrypted_prediction_shorten = list(encrypted_prediction)[:ENCRYPTED_DATA_BROWSER_LIMIT]138 encrypted_prediction_shorten_hex = ''.join(f'{i:02x}' for i in encrypted_prediction_shorten)139 return encrypted_prediction_shorten_hex140 141 142def decrypt_prediction(user_id):143 encoded_data_path = Path(f"tmp/tmp_encrypted_prediction_{user_id}.npy")144 if not user_id:145 raise gr.Error("You need to generate FHE keys first.")146 if not encoded_data_path.is_file():147 raise gr.Error("No encrypted prediction was found. Run the prediction over the encrypted data first.")148 149 # Read encrypted_prediction from the file150 encrypted_prediction = numpy.load(encoded_data_path).tobytes()151 152 fhe_api = FHEModelClient("sentiment_fhe_model/deployment", f".fhe_keys/{user_id}")153 fhe_api.load()154 155 # We need to retrieve the private key that matches the client specs (see issue #18)156 fhe_api.generate_private_and_evaluation_keys(force=False)157 158 predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)159 return {160 "negative": predictions[0][0],161 "neutral": predictions[0][1],162 "positive": predictions[0][2],163 }164 165 166demo = gr.Blocks()167 168 169print("Starting the demo...")170with demo:171 172 gr.Markdown(173 """174<p align="center">175 <img width=200 src="https://user-images.githubusercontent.com/5758427/197816413-d9cddad3-ba38-4793-847d-120975e1da11.png">176</p>177 178<h2 align="center">Machine Learning, Natural Language Processing and Fully Homomorphic Encryption to do Sentiment Analysis on Encrypted data.</h2>179 180<p align="center">181 <a href="https://github.com/zama-ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197972109-faaaff3e-10e2-4ab6-80f5-7531f7cfb08f.png">Concrete-ML</a>182 —183 <a href="https://docs.zama.ai/concrete-ml"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197976802-fddd34c5-f59a-48d0-9bff-7ad1b00cb1fb.png">Documentation</a>184 —185 <a href="https://community.zama.ai"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197977153-8c9c01a7-451a-4993-8e10-5a6ed5343d02.png">Community support forum</a>186 —187 <a href="https://twitter.com/zama_fhe"> <img style="vertical-align: middle; display:inline-block; margin-right: 3px;" width=15 src="https://user-images.githubusercontent.com/5758427/197975044-bab9d199-e120-433b-b3be-abd73b211a54.png">@zama_fhe</a>188</p>189 190<p align="center">191 <img src="https://user-images.githubusercontent.com/7602572/202997494-4ce17b99-9739-4b2c-9f99-e93cca661361.png">192</p>193 194<p align="center">195 <img src="https://user-images.githubusercontent.com/7602572/202998030-5883d817-7b16-406a-8052-b5a0ffe5ec9b.png">196</p>197"""198 )199 200 201 202 # FIXME: make it smaller and in the middle203 # gr.Image("Zama.svg")204 205 gr.Markdown(206 """207 <p align="center">208 </p>209 <p align="center">210 </p>211 """212 )213 214 gr.Markdown("## Notes")215 gr.Markdown(216 """217- The private key is used to encrypt and decrypt the data and shall never be shared.218- The evaluation key is a public key that the server needs to process encrypted data.219"""220 )221 222 gr.Markdown("# Step 1: Generate the keys")223 224 b_gen_key_and_install = gr.Button("Generate the keys and send public part to server")225 226 evaluation_key = gr.Textbox(227 label="Evaluation key (truncated):",228 max_lines=4,229 interactive=False,230 )231 232 user_id = gr.Textbox(233 label="",234 max_lines=4,235 interactive=False,236 visible=False237 )238 239 size_evaluation_key = gr.Number(240 label="Size of the evalution key (in bytes):", value=0, interactive=False241 )242 243 gr.Markdown("# Step 2: Provide a message")244 gr.Markdown("## Client side")245 gr.Markdown(246 "Enter a sensitive text message you received and would like to do sentiment analysis on (ideas: the last text message of your boss.... or lover)."247 )248 text = gr.Textbox(label="Enter a message:", value="I really like your work recently")249 250 gr.Markdown("# Step 3: Encode the message with the private key")251 b_encode_quantize_text = gr.Button(252 "Encode, quantize and encrypt the text with transformer vectorizer, and send to server"253 )254 size_text = gr.Number(label="Size of the text (in bytes):", value="0", interactive=False)255 256 with gr.Row():257 encoding = gr.Textbox(258 label="Transformer representation:",259 max_lines=4,260 interactive=False,261 )262 quantized_encoding = gr.Textbox(263 label="Quantized transformer representation:", max_lines=4, interactive=False264 )265 encrypted_quantized_encoding = gr.Textbox(266 label="Encrypted quantized transformer representation (truncated):",267 max_lines=4,268 interactive=False,269 )270 with gr.Row():271 size_encoding = gr.Number(label="Size (in bytes):", value=0, interactive=False)272 size_quantized_encoding = gr.Number(label="Size (in bytes):", value=0, interactive=False)273 size_encrypted_quantized_encoding = gr.Number(274 label="Size (in bytes):",275 value=0,276 interactive=False,277 )278 279 gr.Markdown("# Step 4: Run the FHE evaluation")280 gr.Markdown("## Server side")281 gr.Markdown(282 "The encrypted value is received by the server. Thanks to the evaluation key and to FHE, the server can compute the (encrypted) prediction directly over encrypted values. Once the computation is finished, the server returns the encrypted prediction to the client."283 )284 285 b_run_fhe = gr.Button("Run FHE execution there")286 encrypted_prediction = gr.Textbox(287 label="Encrypted prediction (truncated):",288 max_lines=4,289 interactive=False,290 )291 292 gr.Markdown("# Step 5: Decrypt the sentiment")293 gr.Markdown("## Client side")294 gr.Markdown(295 "The encrypted sentiment is sent back to client, who can finally decrypt it with its private key. Only the client is aware of the original tweet and the prediction."296 )297 b_decrypt_prediction = gr.Button("Decrypt prediction")298 299 labels_sentiment = gr.Label(label="Sentiment:")300 301 # Button for key generation302 b_gen_key_and_install.click(keygen, inputs=[], outputs=[evaluation_key, size_evaluation_key, user_id])303 304 # Button to quantize and encrypt305 b_encode_quantize_text.click(306 encode_quantize_encrypt,307 inputs=[text, user_id],308 outputs=[309 encoding,310 quantized_encoding,311 encrypted_quantized_encoding,312 size_text,313 size_encoding,314 size_quantized_encoding,315 size_encrypted_quantized_encoding,316 ],317 )318 319 # Button to send the encodings to the server using post at (localhost:8000/predict_sentiment)320 b_run_fhe.click(run_fhe, inputs=[user_id], outputs=[encrypted_prediction])321 322 # Button to decrypt the prediction on the client323 b_decrypt_prediction.click(decrypt_prediction, inputs=[user_id], outputs=[labels_sentiment])324 gr.Markdown(325 "The app was built with [Concrete-ML](https://github.com/zama-ai/concrete-ml), a Privacy-Preserving Machine Learning (PPML) open-source set of tools by [Zama](https://zama.ai/). Try it yourself and don't forget to star on Github ⭐."326 )327demo.launch(share=False)328