{ "cells": [ { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "encoding_width = 2**13 # restricted by ckks size\n", "database_size = 100" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Client Setup" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "from Pyfhel import Pyfhel, PyCtxt" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Client] Initializing Pyfhel session and data...\n" ] } ], "source": [ "print(f\"[Client] Initializing Pyfhel session and data...\")\n", "HE_client = Pyfhel() # Creating empty Pyfhel object\n", "ckks_params = {\n", " 'scheme': 'CKKS', # can also be 'ckks'\n", " 'n': 2**14, # Polynomial modulus degree. For CKKS, n/2 values can be\n", " # encoded in a single ciphertext. \n", " # Typ. 2^D for D in [10, 15]\n", " 'scale': 2**30, # All the encodings will use it for float->fixed point\n", " # conversion: x_fix = round(x_float * scale)\n", " # You can use this as default scale or use a different\n", " # scale on each operation (set in HE.encryptFrac)\n", " 'qi_sizes': [60, 30, 30, 30, 60] # Number of bits of each prime in the chain. \n", " # Intermediate values should be close to log2(scale)\n", " # for each operation, to have small rounding errors.\n", "}\n", "HE_client.contextGen(**ckks_params) # Generate context for bfv scheme\n", "HE_client.keyGen() # Generates both a public and a private key\n", "HE_client.relinKeyGen()\n", "HE_client.rotateKeyGen()" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "# Generate and encrypt query vector\n", "x = np.random.rand(encoding_width)\n", "#cx = np.array([HE_client.encrypt(x[j]) for j in range(len(x))])\n", "# precompute 1 - |x|^2 for query vector\n", "du = 1 - x @ x\n", "cx = HE_client.encrypt(x)" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Client] sending HE_client= and cx=\n" ] } ], "source": [ "# Serializing data and public context information\n", "\n", "s_context = HE_client.to_bytes_context()\n", "s_public_key = HE_client.to_bytes_public_key()\n", "s_relin_key = HE_client.to_bytes_relin_key()\n", "s_rotate_key = HE_client.to_bytes_rotate_key()\n", "#s_cx = [cx[j].to_bytes() for j in range(len(cx))]\n", "s_cx = cx.to_bytes()\n", "\n", "print(f\"[Client] sending HE_client={HE_client} and cx={cx}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Server Mock" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "def hyperbolic_distance_parts(u, v): # returns only the numerator and denominator of the hyperbolic distance formula\n", " diff = u - v\n", " #du = -(1 - u @ u) # for some reason we need to negate this\n", " #dv = -(1 - v @ v) # for some reason we need to negate this\n", " #return diff @ diff, du * dv # returns the numerator and denominator\n", " return diff @ diff # returns the numerator and denominator\n" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [], "source": [ "# document matrix containing rows of document encoding vectors\n", "D = np.random.rand(database_size, encoding_width)\n", "# precompute 1 - |D|^2 for each row vector in D\n", "dv = []\n", "for i in range(len(D)):\n", " v = D[i]\n", " dv.append(1 - v @ v)\n", "dv = np.array(dv)" ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Server] received HE_server= and cx=\n", "[Server] Distances computed! Responding: res=[, , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , ]\n" ] } ], "source": [ "HE_server = Pyfhel()\n", "HE_server.from_bytes_context(s_context)\n", "HE_server.from_bytes_public_key(s_public_key)\n", "HE_server.from_bytes_relin_key(s_relin_key)\n", "HE_server.from_bytes_rotate_key(s_rotate_key)\n", "#cx = np.array([PyCtxt(pyfhel=HE_server, bytestring=s_cx[j]) for j in range(len(s_cx))])\n", "cx = PyCtxt(pyfhel=HE_server, bytestring=s_cx)\n", "print(f\"[Server] received HE_server={HE_server} and cx={cx}\")\n", "\n", "# Encode each document weights in plaintext\n", "res = []\n", "\n", "for i in range(len(D)):\n", " #d = np.array(D[i])\n", " #cd = np.array([HE_server.encrypt(d[j]) for j in range(len(d))])\n", " cd = HE_server.encrypt(D[i])\n", " # Compute distance bewteen recieved query and D[i]\n", " res.append(hyperbolic_distance_parts(cx, cd))\n", "\n", "s_res = [res[j].to_bytes() for j in range(len(res))]\n", "\n", "print(f\"[Server] Distances computed! Responding: res={res}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note that the time is mostly restricted by database size and not encoding size" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Client Parse Response" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "def hyperbolic_distance(u, v):\n", " num = ((u - v) @ (u - v))\n", " den = (1 - (u @ u)) * (1 - (v @ v))\n", " return np.arccosh(1 + 2 * (num / den))" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "#res = np.array([HE_client.decrypt(c_res[j]) for j in range(len(c_res))])[:,0]\n", "#res = HE_client.decrypt(c_res)\n", "c_res = []\n", "for i in range(len(s_res)):\n", " c_num = PyCtxt(pyfhel=HE_server, bytestring=s_res[i])\n", " #c_den = PyCtxt(pyfhel=HE_server, bytestring=s_res[i][1])\n", " p_num = HE_client.decrypt(c_num)[0]\n", " #p_den = HE_client.decrypt(c_den)[0]\n", " #dist = np.arccosh(1 + 2 * (p_num / p_den))\n", " # compute final score\n", " dist = np.arccosh(1 + 2 * (p_num / (du * dv[i])))\n", " #print(dist)\n", " c_res.append(dist)\n", "\n", "# Checking result\n", "expected_res = [hyperbolic_distance(x, np.array(w)) for w in D]\n", "#print(f\"[Client] Response received! \\nResult is {c_res} \\nShould be {expected}\\nDiff {np.abs(np.array(c_res) - np.array(expected))}\")\n", "for i in range(len(c_res)):\n", " result = c_res[i]\n", " expected = expected_res[i]\n", " if np.abs(result - expected) < 1e-3:\n", " pass\n", " else:\n", " print(f\"got: {result}, expected: {expected}\")\n", " assert False" ] } ], "metadata": { "kernelspec": { "display_name": "cs239", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.20" } }, "nbformat": 4, "nbformat_minor": 2 }