ppaihack/CipherClause
0
1import gradio as gr2from requests import head3from transformer_vectorizer import TransformerVectorizer4from sklearn.feature_extraction.text import TfidfVectorizer5import numpy as np 6from concrete.ml.deployment import FHEModelClient7import numpy8import os9from pathlib import Path10import requests11import json12import base6413import subprocess14import shutil15import time16import easyocr17import PyPDF218import os19 20reader = easyocr.Reader(['en']) 21# This repository's directory22REPO_DIR = Path(__file__).parent23 24subprocess.Popen(["uvicorn", "server:app"], cwd=REPO_DIR)25 26# Wait 5 sec for the server to start27time.sleep(5)28 29# Encrypted data limit for the browser to display30# (encrypted data is too large to display in the browser)31ENCRYPTED_DATA_BROWSER_LIMIT = 50032N_USER_KEY_STORED = 2033model_names=['financial_rating','legal_rating']34 35 36FHE_MODEL_PATH = "deployment/financial_rating"37FHE_LEGAL_PATH = "deployment/legal_rating"38#FHE_LEGAL_PATH="deployment/legal_rating"39 40print("Loading the transformer model...")41 42# Initialize the transformer vectorizer43transformer_vectorizer = TransformerVectorizer()44vectorizer = TfidfVectorizer()45def process_input(input_type, user_input, uploaded_file):46 print('ooooocr')47 if input_type == "File Upload" and uploaded_file is not None:48 # 读取上传的文件49 with open(uploaded_file.name, "rb") as f:50 image = f.read()51 results = reader.readtext(image)52 # 提取识别的文本53 extracted_text = ' '.join([text[1] for text in results])54 print("提取的文本:")55 print(extracted_text)56 return extracted_text57 elif input_type == "Text Input":58 return user_input59''' 60def process_input(input_type, user_input, uploaded_file):61 if input_type == "File Upload" and uploaded_file is not None:62 file_ext = os.path.splitext(uploaded_file.name)[1].lower()63 extracted_text = ""64 65 if file_ext in ['.jpg', '.jpeg', '.png']:66 # 处理图片文件67 results = reader.readtext(uploaded_file.name)68 extracted_text = ' '.join([text[1] for text in results])69 print("从图片提取的文本:")70 print(extracted_text)71 72 elif file_ext == '.txt':73 # 处理TXT文件74 with open(uploaded_file.name, 'r', encoding='utf-8') as f:75 extracted_text = f.read()76 print("从TXT文件提取的文本:")77 print(extracted_text)78 79 elif file_ext == '.pdf':80 # 处理PDF文件81 with open(uploaded_file.name, 'rb') as f:82 reader_pdf = PyPDF2.PdfReader(f)83 for page_num in range(len(reader_pdf.pages)):84 page = reader_pdf.pages[page_num]85 extracted_text += page.extract_text() + "\n"86 print("从PDF文件提取的文本:")87 print(extracted_text)88 89 else:90 return "不支持的文件类型。请上传 .jpg, .jpeg, .png, .txt 或 .pdf 文件。"91 92 return extracted_text93 94 elif input_type == "Text Input":95 return user_input96 else:97 return "无效的输入类型或未上传文件。"98'''99 100def toggle_visibility(input_type):101 user_input_visible = input_type == "Text Input"102 file_upload_visible = input_type == "File Upload"103 return gr.update(visible=user_input_visible), gr.update(visible=file_upload_visible)104 105 106def clean_tmp_directory():107 # Allow 20 user keys to be stored.108 # Once that limitation is reached, deleted the oldest.109 path_sub_directories = sorted([f for f in Path(".fhe_keys/").iterdir() if f.is_dir()], key=os.path.getmtime)110 111 user_ids = []112 if len(path_sub_directories) > N_USER_KEY_STORED:113 n_files_to_delete = len(path_sub_directories) - N_USER_KEY_STORED114 for p in path_sub_directories[:n_files_to_delete]:115 user_ids.append(p.name)116 shutil.rmtree(p)117 118 list_files_tmp = Path("tmp/").iterdir()119 # Delete all files related to user_id120 for file in list_files_tmp:121 for user_id in user_ids:122 if file.name.endswith(f"{user_id}.npy"):123 file.unlink()124mes=[]125 126def keygen(selected_tasks):127 # Clean tmp directory if needed128 clean_tmp_directory()129 130 print("Initializing FHEModelClient...")131 132 133 134 if not selected_tasks:135 return "choose a task first" # 修改提示信息为英文136 user_id = numpy.random.randint(0, 2**32)137 if "legal_rating" in selected_tasks:138 model_names.append('legal_rating')139 # Let's create a user_id140 141 fhe_api= FHEModelClient(FHE_LEGAL_PATH, f".fhe_keys/{user_id}")142 143 144 if "financial_rating" in selected_tasks:145 model_names.append('financial_rating')146 147 fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")148 149 # Let's create a user_id150 151 152 fhe_api.load()153 154 155 # Generate a fresh key156 fhe_api.generate_private_and_evaluation_keys(force=True)157 evaluation_key = fhe_api.get_serialized_evaluation_keys()158 159 # Save evaluation_key in a file, since too large to pass through regular Gradio160 # buttons, https://github.com/gradio-app/gradio/issues/1877161 numpy.save(f"tmp/tmp_evaluation_key_{user_id}.npy", evaluation_key)162 163 return [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], user_id]164 165 166 167 168 169def encode_quantize_encrypt(text, user_id):170 if not user_id:171 raise gr.Error("You need to generate FHE keys first.")172 if "legal_rating" in model_names:173 fhe_api = FHEModelClient(FHE_LEGAL_PATH, f".fhe_keys/{user_id}")174 encodings =vectorizer.fit_transform([text]).toarray()175 if encodings.shape[1] < 1736:176 # 在后面填充零177 padding = np.zeros((1, 1736 - encodings.shape[1]))178 encodings = np.hstack((encodings, padding))179 elif encodings.shape[1] > 1736:180 # 截取前1736列181 encodings = encodings[:, :1736]182 else:183 fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")184 encodings = transformer_vectorizer.transform([text])185 186 fhe_api.load()187 quantized_encodings = fhe_api.model.quantize_input(encodings).astype(numpy.uint8)188 encrypted_quantized_encoding = fhe_api.quantize_encrypt_serialize(encodings)189 190 # Save encrypted_quantized_encoding in a file, since too large to pass through regular Gradio191 # buttons, https://github.com/gradio-app/gradio/issues/1877192 numpy.save(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy", encrypted_quantized_encoding)193 194 # Compute size195 encrypted_quantized_encoding_shorten = list(encrypted_quantized_encoding)[:ENCRYPTED_DATA_BROWSER_LIMIT]196 encrypted_quantized_encoding_shorten_hex = ''.join(f'{i:02x}' for i in encrypted_quantized_encoding_shorten)197 return (198 encodings[0],199 quantized_encodings[0],200 encrypted_quantized_encoding_shorten_hex,201 )202 203 204 205def run_fhe(user_id):206 encoded_data_path = Path(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy")207 if not user_id:208 raise gr.Error("You need to generate FHE keys first.")209 if not encoded_data_path.is_file():210 raise gr.Error("No encrypted data was found. Encrypt the data before trying to predict.")211 212 # Read encrypted_quantized_encoding from the file213 encrypted_quantized_encoding = numpy.load(encoded_data_path)214 215 # Read evaluation_key from the file216 evaluation_key = numpy.load(f"tmp/tmp_evaluation_key_{user_id}.npy")217 218 # Use base64 to encode the encodings and evaluation key219 encrypted_quantized_encoding = base64.b64encode(encrypted_quantized_encoding).decode()220 encoded_evaluation_key = base64.b64encode(evaluation_key).decode()221 222 query = {}223 query["evaluation_key"] = encoded_evaluation_key224 query["encrypted_encoding"] = encrypted_quantized_encoding225 headers = {"Content-type": "application/json"}226 if "legal_rating" in model_names:227 response = requests.post(228 "http://localhost:8000/predict_legal", data=json.dumps(query), headers=headers229 )230 else:231 response = requests.post(232 "http://localhost:8000/predict_sentiment", data=json.dumps(query), headers=headers233 )234 encrypted_prediction = base64.b64decode(response.json()["encrypted_prediction"])235 236 # Save encrypted_prediction in a file, since too large to pass through regular Gradio237 # buttons, https://github.com/gradio-app/gradio/issues/1877238 numpy.save(f"tmp/tmp_encrypted_prediction_{user_id}.npy", encrypted_prediction)239 encrypted_prediction_shorten = list(encrypted_prediction)[:ENCRYPTED_DATA_BROWSER_LIMIT]240 encrypted_prediction_shorten_hex = ''.join(f'{i:02x}' for i in encrypted_prediction_shorten)241 return encrypted_prediction_shorten_hex242 243 244def decrypt_prediction(user_id):245 encoded_data_path = Path(f"tmp/tmp_encrypted_prediction_{user_id}.npy")246 if not user_id:247 raise gr.Error("You need to generate FHE keys first.")248 if not encoded_data_path.is_file():249 raise gr.Error("No encrypted prediction was found. Run the prediction over the encrypted data first.")250 251 # Read encrypted_prediction from the file252 encrypted_prediction = numpy.load(encoded_data_path).tobytes()253 254 if "legal_rating" in model_names:255 fhe_api = FHEModelClient(FHE_LEGAL_PATH, f".fhe_keys/{user_id}")256 257 fhe_api = FHEModelClient(FHE_MODEL_PATH, f".fhe_keys/{user_id}")258 fhe_api.load()259 260 # We need to retrieve the private key that matches the client specs (see issue #18)261 fhe_api.generate_private_and_evaluation_keys(force=False)262 263 predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)264 print(predictions)265 266 return {267 "low_relative": predictions[0][0],268 "medium_relative": predictions[0][1],269 "high_relative": predictions[0][2],270 }271 272 273demo = gr.Blocks()274 275 276print("Starting the demo...")277with demo:278 279 gr.Markdown(280 """281 282<h2 align="center">📄Cipher Clause</h2>283 <p align="center">284 <img width=600 src="https://www.helloimg.com/i/2024/09/28/66f7f6701bcfb.jpeg">285 </p>286 287"""288 )289 290 291 gr.Markdown(292 """293 <p align="center">294 </p>295 <p align="center">296 </p>297 """298 )299 300 gr.Markdown("## Notes")301 gr.Markdown(302 """303- The private key is used to encrypt and decrypt the data and shall never be shared.304- The evaluation key is a public key that the server needs to process encrypted data.305"""306 )307 gr.Markdown(308 """309 <hr/>310 """311 )312 gr.Markdown("# Step 0: Select Task") 313 task_checkbox = gr.CheckboxGroup(314 choices=["legal_rating", "financial_rating"],315 label="select_tasks"316 )317 gr.Markdown(318 """319 <hr/>320 """321 )322 gr.Markdown("# Step 1: Generate the keys")323 324 b_gen_key_and_install = gr.Button("Generate all the keys and send public part to server")325 326 evaluation_key = gr.Textbox(327 label="Evaluation key (truncated):",328 max_lines=4,329 interactive=False,330 )331 332 user_id = gr.Textbox(333 label="",334 max_lines=4,335 interactive=False,336 visible=False337 )338 gr.Markdown(339 """340<hr/>341 """342 )343 gr.Markdown("# Step 2: Provide a contract or clause")344 gr.Markdown("## Client side")345 gr.Markdown(346 "Enter a contract or clause you want to analysis)."347 )348 input_type = gr.Radio(choices=["Text Input", "File Upload"], label="Select Input Method")349 user_input = gr.Textbox(label="Enter some words:",visible=False,value="The Employee is entitled to two weeks of paid vacation annually, to be scheduled at the mutual convenience of the Employee and Employer.")350 351 file_upload = gr.File(label="Upload File", file_types=[".jpg", ".png",".txt"], visible=False) # Initially hidden352 input_type.change(toggle_visibility, inputs=input_type, outputs=[user_input, file_upload])353 354 submit_button = gr.Button("Submit")355 text = gr.Textbox(label="Extracted Text")356 357 submit_button.click(process_input, inputs=[input_type, user_input, file_upload], outputs=text)358 359 gr.Markdown(360 """361<hr/>362 """363 )364 gr.Markdown("# Step 3: Encode the message with the private key")365 b_encode_quantize_text = gr.Button(366 "Encode, quantize and encrypt the text with vectorizer, and send to server"367 )368 369 with gr.Row():370 encoding = gr.Textbox(371 label="Representation:",372 max_lines=4,373 interactive=False,374 )375 quantized_encoding = gr.Textbox(376 label="Quantized representation:", max_lines=4, interactive=False377 )378 encrypted_quantized_encoding = gr.Textbox(379 label="Encrypted quantized representation (truncated):",380 max_lines=4,381 interactive=False,382 )383 gr.Markdown(384 """385<hr/>386 """387 )388 gr.Markdown("# Step 4: Run the FHE evaluation")389 gr.Markdown("## Server side")390 gr.Markdown(391 "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."392 )393 394 b_run_fhe = gr.Button("Run FHE execution there")395 encrypted_prediction = gr.Textbox(396 label="Encrypted prediction (truncated):",397 max_lines=4,398 interactive=False,399 )400 gr.Markdown(401 """402<hr/>403 """404 )405 gr.Markdown("# Step 5: Decrypt the class")406 gr.Markdown("## Client side")407 gr.Markdown(408 "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."409 )410 b_decrypt_prediction = gr.Button("Decrypt prediction")411 412 labels_sentiment = gr.Label(label="level:")413 414 # Button for key generation415 b_gen_key_and_install.click(keygen, inputs=[task_checkbox], outputs=[evaluation_key, user_id])416 417 # Button to quantize and encrypt418 b_encode_quantize_text.click(419 encode_quantize_encrypt,420 inputs=[text, user_id],421 outputs=[422 encoding,423 quantized_encoding,424 encrypted_quantized_encoding,425 ],426 )427 428 # Button to send the encodings to the server using post at (localhost:8000/predict_sentiment)429 b_run_fhe.click(run_fhe, inputs=[user_id], outputs=[encrypted_prediction])430 431 # Button to decrypt the prediction on the client432 b_decrypt_prediction.click(decrypt_prediction, inputs=[user_id], outputs=[labels_sentiment])433 gr.Markdown(434 "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 ⭐."435 )436demo.launch(share=False)