wm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "import watermarking.watermark_processor\n",
    "\n",
    "watermarking.watermark_processor = reload(watermarking.watermark_processor)\n",
    "from watermarking.watermark_processor import RepetitionPenaltyLogitsProcessor\n",
    "from transformers import LogitsProcessorList, MinLengthLogitsProcessor, LogitsProcessor\n",
    "from watermarking.utils.text_tools import truncate\n",
    "from watermarking.utils.load_local import load_local_model_or_tokenizer\n",
    "from datasets import load_dataset, load_from_disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading opt-1.3b tokenizer from /home/wanglean/llm-ckpts/tokenizers/opt-1.3b ...\n",
      "finished loading\n",
      "loading opt-1.3b model from /home/wanglean/llm-ckpts/models/opt-1.3b ...\n",
      "finished loading\n"
     ]
    }
   ],
   "source": [
    "tokenizer = load_local_model_or_tokenizer(\"facebook/opt-1.3b\", 'tokenizer')\n",
    "model = load_local_model_or_tokenizer(\"facebook/opt-1.3b\", 'model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 466,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "ename": "SSLError",
     "evalue": "HTTPSConnectionPool(host='huggingface.co', port=443): Max retries exceeded with url: /gpt2/resolve/main/tokenizer_config.json (Caused by SSLError(SSLEOFError(8, 'EOF occurred in violation of protocol (_ssl.c:997)')))",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mSSLEOFError\u001B[0m                               Traceback (most recent call last)",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/connectionpool.py:700\u001B[0m, in \u001B[0;36mHTTPConnectionPool.urlopen\u001B[0;34m(self, method, url, body, headers, retries, redirect, assert_same_host, timeout, pool_timeout, release_conn, chunked, body_pos, **response_kw)\u001B[0m\n\u001B[1;32m    699\u001B[0m \u001B[39mif\u001B[39;00m is_new_proxy_conn \u001B[39mand\u001B[39;00m http_tunnel_required:\n\u001B[0;32m--> 700\u001B[0m     \u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49m_prepare_proxy(conn)\n\u001B[1;32m    702\u001B[0m \u001B[39m# Make the request on the httplib connection object.\u001B[39;00m\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/connectionpool.py:996\u001B[0m, in \u001B[0;36mHTTPSConnectionPool._prepare_proxy\u001B[0;34m(self, conn)\u001B[0m\n\u001B[1;32m    994\u001B[0m     conn\u001B[39m.\u001B[39mtls_in_tls_required \u001B[39m=\u001B[39m \u001B[39mTrue\u001B[39;00m\n\u001B[0;32m--> 996\u001B[0m conn\u001B[39m.\u001B[39;49mconnect()\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/connection.py:414\u001B[0m, in \u001B[0;36mHTTPSConnection.connect\u001B[0;34m(self)\u001B[0m\n\u001B[1;32m    412\u001B[0m     context\u001B[39m.\u001B[39mload_default_certs()\n\u001B[0;32m--> 414\u001B[0m \u001B[39mself\u001B[39m\u001B[39m.\u001B[39msock \u001B[39m=\u001B[39m ssl_wrap_socket(\n\u001B[1;32m    415\u001B[0m     sock\u001B[39m=\u001B[39;49mconn,\n\u001B[1;32m    416\u001B[0m     keyfile\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mkey_file,\n\u001B[1;32m    417\u001B[0m     certfile\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mcert_file,\n\u001B[1;32m    418\u001B[0m     key_password\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mkey_password,\n\u001B[1;32m    419\u001B[0m     ca_certs\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mca_certs,\n\u001B[1;32m    420\u001B[0m     ca_cert_dir\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mca_cert_dir,\n\u001B[1;32m    421\u001B[0m     ca_cert_data\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mca_cert_data,\n\u001B[1;32m    422\u001B[0m     server_hostname\u001B[39m=\u001B[39;49mserver_hostname,\n\u001B[1;32m    423\u001B[0m     ssl_context\u001B[39m=\u001B[39;49mcontext,\n\u001B[1;32m    424\u001B[0m     tls_in_tls\u001B[39m=\u001B[39;49mtls_in_tls,\n\u001B[1;32m    425\u001B[0m )\n\u001B[1;32m    427\u001B[0m \u001B[39m# If we're using all defaults and the connection\u001B[39;00m\n\u001B[1;32m    428\u001B[0m \u001B[39m# is TLSv1 or TLSv1.1 we throw a DeprecationWarning\u001B[39;00m\n\u001B[1;32m    429\u001B[0m \u001B[39m# for the host.\u001B[39;00m\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/util/ssl_.py:449\u001B[0m, in \u001B[0;36mssl_wrap_socket\u001B[0;34m(sock, keyfile, certfile, cert_reqs, ca_certs, server_hostname, ssl_version, ciphers, ssl_context, ca_cert_dir, key_password, ca_cert_data, tls_in_tls)\u001B[0m\n\u001B[1;32m    448\u001B[0m \u001B[39mif\u001B[39;00m send_sni:\n\u001B[0;32m--> 449\u001B[0m     ssl_sock \u001B[39m=\u001B[39m _ssl_wrap_socket_impl(\n\u001B[1;32m    450\u001B[0m         sock, context, tls_in_tls, server_hostname\u001B[39m=\u001B[39;49mserver_hostname\n\u001B[1;32m    451\u001B[0m     )\n\u001B[1;32m    452\u001B[0m \u001B[39melse\u001B[39;00m:\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/util/ssl_.py:493\u001B[0m, in \u001B[0;36m_ssl_wrap_socket_impl\u001B[0;34m(sock, ssl_context, tls_in_tls, server_hostname)\u001B[0m\n\u001B[1;32m    492\u001B[0m \u001B[39mif\u001B[39;00m server_hostname:\n\u001B[0;32m--> 493\u001B[0m     \u001B[39mreturn\u001B[39;00m ssl_context\u001B[39m.\u001B[39;49mwrap_socket(sock, server_hostname\u001B[39m=\u001B[39;49mserver_hostname)\n\u001B[1;32m    494\u001B[0m \u001B[39melse\u001B[39;00m:\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/ssl.py:513\u001B[0m, in \u001B[0;36mSSLContext.wrap_socket\u001B[0;34m(self, sock, server_side, do_handshake_on_connect, suppress_ragged_eofs, server_hostname, session)\u001B[0m\n\u001B[1;32m    507\u001B[0m \u001B[39mdef\u001B[39;00m \u001B[39mwrap_socket\u001B[39m(\u001B[39mself\u001B[39m, sock, server_side\u001B[39m=\u001B[39m\u001B[39mFalse\u001B[39;00m,\n\u001B[1;32m    508\u001B[0m                 do_handshake_on_connect\u001B[39m=\u001B[39m\u001B[39mTrue\u001B[39;00m,\n\u001B[1;32m    509\u001B[0m                 suppress_ragged_eofs\u001B[39m=\u001B[39m\u001B[39mTrue\u001B[39;00m,\n\u001B[1;32m    510\u001B[0m                 server_hostname\u001B[39m=\u001B[39m\u001B[39mNone\u001B[39;00m, session\u001B[39m=\u001B[39m\u001B[39mNone\u001B[39;00m):\n\u001B[1;32m    511\u001B[0m     \u001B[39m# SSLSocket class handles server_hostname encoding before it calls\u001B[39;00m\n\u001B[1;32m    512\u001B[0m     \u001B[39m# ctx._wrap_socket()\u001B[39;00m\n\u001B[0;32m--> 513\u001B[0m     \u001B[39mreturn\u001B[39;00m \u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49msslsocket_class\u001B[39m.\u001B[39;49m_create(\n\u001B[1;32m    514\u001B[0m         sock\u001B[39m=\u001B[39;49msock,\n\u001B[1;32m    515\u001B[0m         server_side\u001B[39m=\u001B[39;49mserver_side,\n\u001B[1;32m    516\u001B[0m         do_handshake_on_connect\u001B[39m=\u001B[39;49mdo_handshake_on_connect,\n\u001B[1;32m    517\u001B[0m         suppress_ragged_eofs\u001B[39m=\u001B[39;49msuppress_ragged_eofs,\n\u001B[1;32m    518\u001B[0m         server_hostname\u001B[39m=\u001B[39;49mserver_hostname,\n\u001B[1;32m    519\u001B[0m         context\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m,\n\u001B[1;32m    520\u001B[0m         session\u001B[39m=\u001B[39;49msession\n\u001B[1;32m    521\u001B[0m     )\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/ssl.py:1071\u001B[0m, in \u001B[0;36mSSLSocket._create\u001B[0;34m(cls, sock, server_side, do_handshake_on_connect, suppress_ragged_eofs, server_hostname, context, session)\u001B[0m\n\u001B[1;32m   1070\u001B[0m             \u001B[39mraise\u001B[39;00m \u001B[39mValueError\u001B[39;00m(\u001B[39m\"\u001B[39m\u001B[39mdo_handshake_on_connect should not be specified for non-blocking sockets\u001B[39m\u001B[39m\"\u001B[39m)\n\u001B[0;32m-> 1071\u001B[0m         \u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mdo_handshake()\n\u001B[1;32m   1072\u001B[0m \u001B[39mexcept\u001B[39;00m (\u001B[39mOSError\u001B[39;00m, \u001B[39mValueError\u001B[39;00m):\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/ssl.py:1342\u001B[0m, in \u001B[0;36mSSLSocket.do_handshake\u001B[0;34m(self, block)\u001B[0m\n\u001B[1;32m   1341\u001B[0m         \u001B[39mself\u001B[39m\u001B[39m.\u001B[39msettimeout(\u001B[39mNone\u001B[39;00m)\n\u001B[0;32m-> 1342\u001B[0m     \u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49m_sslobj\u001B[39m.\u001B[39;49mdo_handshake()\n\u001B[1;32m   1343\u001B[0m \u001B[39mfinally\u001B[39;00m:\n",
      "\u001B[0;31mSSLEOFError\u001B[0m: EOF occurred in violation of protocol (_ssl.c:997)",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001B[0;31mMaxRetryError\u001B[0m                             Traceback (most recent call last)",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/requests/adapters.py:489\u001B[0m, in \u001B[0;36mHTTPAdapter.send\u001B[0;34m(self, request, stream, timeout, verify, cert, proxies)\u001B[0m\n\u001B[1;32m    488\u001B[0m \u001B[39mif\u001B[39;00m \u001B[39mnot\u001B[39;00m chunked:\n\u001B[0;32m--> 489\u001B[0m     resp \u001B[39m=\u001B[39m conn\u001B[39m.\u001B[39;49murlopen(\n\u001B[1;32m    490\u001B[0m         method\u001B[39m=\u001B[39;49mrequest\u001B[39m.\u001B[39;49mmethod,\n\u001B[1;32m    491\u001B[0m         url\u001B[39m=\u001B[39;49murl,\n\u001B[1;32m    492\u001B[0m         body\u001B[39m=\u001B[39;49mrequest\u001B[39m.\u001B[39;49mbody,\n\u001B[1;32m    493\u001B[0m         headers\u001B[39m=\u001B[39;49mrequest\u001B[39m.\u001B[39;49mheaders,\n\u001B[1;32m    494\u001B[0m         redirect\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    495\u001B[0m         assert_same_host\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    496\u001B[0m         preload_content\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    497\u001B[0m         decode_content\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    498\u001B[0m         retries\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49mmax_retries,\n\u001B[1;32m    499\u001B[0m         timeout\u001B[39m=\u001B[39;49mtimeout,\n\u001B[1;32m    500\u001B[0m     )\n\u001B[1;32m    502\u001B[0m \u001B[39m# Send the request.\u001B[39;00m\n\u001B[1;32m    503\u001B[0m \u001B[39melse\u001B[39;00m:\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/connectionpool.py:787\u001B[0m, in \u001B[0;36mHTTPConnectionPool.urlopen\u001B[0;34m(self, method, url, body, headers, retries, redirect, assert_same_host, timeout, pool_timeout, release_conn, chunked, body_pos, **response_kw)\u001B[0m\n\u001B[1;32m    785\u001B[0m     e \u001B[39m=\u001B[39m ProtocolError(\u001B[39m\"\u001B[39m\u001B[39mConnection aborted.\u001B[39m\u001B[39m\"\u001B[39m, e)\n\u001B[0;32m--> 787\u001B[0m retries \u001B[39m=\u001B[39m retries\u001B[39m.\u001B[39;49mincrement(\n\u001B[1;32m    788\u001B[0m     method, url, error\u001B[39m=\u001B[39;49me, _pool\u001B[39m=\u001B[39;49m\u001B[39mself\u001B[39;49m, _stacktrace\u001B[39m=\u001B[39;49msys\u001B[39m.\u001B[39;49mexc_info()[\u001B[39m2\u001B[39;49m]\n\u001B[1;32m    789\u001B[0m )\n\u001B[1;32m    790\u001B[0m retries\u001B[39m.\u001B[39msleep()\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/urllib3/util/retry.py:592\u001B[0m, in \u001B[0;36mRetry.increment\u001B[0;34m(self, method, url, response, error, _pool, _stacktrace)\u001B[0m\n\u001B[1;32m    591\u001B[0m \u001B[39mif\u001B[39;00m new_retry\u001B[39m.\u001B[39mis_exhausted():\n\u001B[0;32m--> 592\u001B[0m     \u001B[39mraise\u001B[39;00m MaxRetryError(_pool, url, error \u001B[39mor\u001B[39;00m ResponseError(cause))\n\u001B[1;32m    594\u001B[0m log\u001B[39m.\u001B[39mdebug(\u001B[39m\"\u001B[39m\u001B[39mIncremented Retry for (url=\u001B[39m\u001B[39m'\u001B[39m\u001B[39m%s\u001B[39;00m\u001B[39m'\u001B[39m\u001B[39m): \u001B[39m\u001B[39m%r\u001B[39;00m\u001B[39m\"\u001B[39m, url, new_retry)\n",
      "\u001B[0;31mMaxRetryError\u001B[0m: HTTPSConnectionPool(host='huggingface.co', port=443): Max retries exceeded with url: /gpt2/resolve/main/tokenizer_config.json (Caused by SSLError(SSLEOFError(8, 'EOF occurred in violation of protocol (_ssl.c:997)')))",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001B[0;31mSSLError\u001B[0m                                  Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[466], line 5\u001B[0m\n\u001B[1;32m      2\u001B[0m os\u001B[38;5;241m.\u001B[39menviron[\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mCURL_CA_BUNDLE\u001B[39m\u001B[38;5;124m'\u001B[39m] \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124m'\u001B[39m\n\u001B[1;32m      4\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mtransformers\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m AutoTokenizer, AutoModelForCausalLM\n\u001B[0;32m----> 5\u001B[0m lm_tokenizer \u001B[38;5;241m=\u001B[39m \u001B[43mAutoTokenizer\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mfrom_pretrained\u001B[49m\u001B[43m(\u001B[49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[38;5;124;43mgpt2\u001B[39;49m\u001B[38;5;124;43m'\u001B[39;49m\u001B[43m)\u001B[49m\n\u001B[1;32m      6\u001B[0m lm_model \u001B[38;5;241m=\u001B[39m AutoModelForCausalLM\u001B[38;5;241m.\u001B[39mfrom_pretrained(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mgpt2\u001B[39m\u001B[38;5;124m'\u001B[39m)\u001B[38;5;241m.\u001B[39mto(\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mcuda:0\u001B[39m\u001B[38;5;124m'\u001B[39m)\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py:634\u001B[0m, in \u001B[0;36mAutoTokenizer.from_pretrained\u001B[0;34m(cls, pretrained_model_name_or_path, *inputs, **kwargs)\u001B[0m\n\u001B[1;32m    631\u001B[0m     \u001B[39mreturn\u001B[39;00m tokenizer_class\u001B[39m.\u001B[39mfrom_pretrained(pretrained_model_name_or_path, \u001B[39m*\u001B[39minputs, \u001B[39m*\u001B[39m\u001B[39m*\u001B[39mkwargs)\n\u001B[1;32m    633\u001B[0m \u001B[39m# Next, let's try to use the tokenizer_config file to get the tokenizer class.\u001B[39;00m\n\u001B[0;32m--> 634\u001B[0m tokenizer_config \u001B[39m=\u001B[39m get_tokenizer_config(pretrained_model_name_or_path, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs)\n\u001B[1;32m    635\u001B[0m \u001B[39mif\u001B[39;00m \u001B[39m\"\u001B[39m\u001B[39m_commit_hash\u001B[39m\u001B[39m\"\u001B[39m \u001B[39min\u001B[39;00m tokenizer_config:\n\u001B[1;32m    636\u001B[0m     kwargs[\u001B[39m\"\u001B[39m\u001B[39m_commit_hash\u001B[39m\u001B[39m\"\u001B[39m] \u001B[39m=\u001B[39m tokenizer_config[\u001B[39m\"\u001B[39m\u001B[39m_commit_hash\u001B[39m\u001B[39m\"\u001B[39m]\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py:478\u001B[0m, in \u001B[0;36mget_tokenizer_config\u001B[0;34m(pretrained_model_name_or_path, cache_dir, force_download, resume_download, proxies, use_auth_token, revision, local_files_only, subfolder, **kwargs)\u001B[0m\n\u001B[1;32m    416\u001B[0m \u001B[39m\"\"\"\u001B[39;00m\n\u001B[1;32m    417\u001B[0m \u001B[39mLoads the tokenizer configuration from a pretrained model tokenizer configuration.\u001B[39;00m\n\u001B[1;32m    418\u001B[0m \n\u001B[0;32m   (...)\u001B[0m\n\u001B[1;32m    475\u001B[0m \u001B[39mtokenizer_config = get_tokenizer_config(\"tokenizer-test\")\u001B[39;00m\n\u001B[1;32m    476\u001B[0m \u001B[39m```\"\"\"\u001B[39;00m\n\u001B[1;32m    477\u001B[0m commit_hash \u001B[39m=\u001B[39m kwargs\u001B[39m.\u001B[39mget(\u001B[39m\"\u001B[39m\u001B[39m_commit_hash\u001B[39m\u001B[39m\"\u001B[39m, \u001B[39mNone\u001B[39;00m)\n\u001B[0;32m--> 478\u001B[0m resolved_config_file \u001B[39m=\u001B[39m cached_file(\n\u001B[1;32m    479\u001B[0m     pretrained_model_name_or_path,\n\u001B[1;32m    480\u001B[0m     TOKENIZER_CONFIG_FILE,\n\u001B[1;32m    481\u001B[0m     cache_dir\u001B[39m=\u001B[39;49mcache_dir,\n\u001B[1;32m    482\u001B[0m     force_download\u001B[39m=\u001B[39;49mforce_download,\n\u001B[1;32m    483\u001B[0m     resume_download\u001B[39m=\u001B[39;49mresume_download,\n\u001B[1;32m    484\u001B[0m     proxies\u001B[39m=\u001B[39;49mproxies,\n\u001B[1;32m    485\u001B[0m     use_auth_token\u001B[39m=\u001B[39;49muse_auth_token,\n\u001B[1;32m    486\u001B[0m     revision\u001B[39m=\u001B[39;49mrevision,\n\u001B[1;32m    487\u001B[0m     local_files_only\u001B[39m=\u001B[39;49mlocal_files_only,\n\u001B[1;32m    488\u001B[0m     subfolder\u001B[39m=\u001B[39;49msubfolder,\n\u001B[1;32m    489\u001B[0m     _raise_exceptions_for_missing_entries\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    490\u001B[0m     _raise_exceptions_for_connection_errors\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    491\u001B[0m     _commit_hash\u001B[39m=\u001B[39;49mcommit_hash,\n\u001B[1;32m    492\u001B[0m )\n\u001B[1;32m    493\u001B[0m \u001B[39mif\u001B[39;00m resolved_config_file \u001B[39mis\u001B[39;00m \u001B[39mNone\u001B[39;00m:\n\u001B[1;32m    494\u001B[0m     logger\u001B[39m.\u001B[39minfo(\u001B[39m\"\u001B[39m\u001B[39mCould not locate the tokenizer configuration file, will try to use the model config instead.\u001B[39m\u001B[39m\"\u001B[39m)\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/transformers/utils/hub.py:409\u001B[0m, in \u001B[0;36mcached_file\u001B[0;34m(path_or_repo_id, filename, cache_dir, force_download, resume_download, proxies, use_auth_token, revision, local_files_only, subfolder, user_agent, _raise_exceptions_for_missing_entries, _raise_exceptions_for_connection_errors, _commit_hash)\u001B[0m\n\u001B[1;32m    406\u001B[0m user_agent \u001B[39m=\u001B[39m http_user_agent(user_agent)\n\u001B[1;32m    407\u001B[0m \u001B[39mtry\u001B[39;00m:\n\u001B[1;32m    408\u001B[0m     \u001B[39m# Load from URL or cache if already cached\u001B[39;00m\n\u001B[0;32m--> 409\u001B[0m     resolved_file \u001B[39m=\u001B[39m hf_hub_download(\n\u001B[1;32m    410\u001B[0m         path_or_repo_id,\n\u001B[1;32m    411\u001B[0m         filename,\n\u001B[1;32m    412\u001B[0m         subfolder\u001B[39m=\u001B[39;49m\u001B[39mNone\u001B[39;49;00m \u001B[39mif\u001B[39;49;00m \u001B[39mlen\u001B[39;49m(subfolder) \u001B[39m==\u001B[39;49m \u001B[39m0\u001B[39;49m \u001B[39melse\u001B[39;49;00m subfolder,\n\u001B[1;32m    413\u001B[0m         revision\u001B[39m=\u001B[39;49mrevision,\n\u001B[1;32m    414\u001B[0m         cache_dir\u001B[39m=\u001B[39;49mcache_dir,\n\u001B[1;32m    415\u001B[0m         user_agent\u001B[39m=\u001B[39;49muser_agent,\n\u001B[1;32m    416\u001B[0m         force_download\u001B[39m=\u001B[39;49mforce_download,\n\u001B[1;32m    417\u001B[0m         proxies\u001B[39m=\u001B[39;49mproxies,\n\u001B[1;32m    418\u001B[0m         resume_download\u001B[39m=\u001B[39;49mresume_download,\n\u001B[1;32m    419\u001B[0m         use_auth_token\u001B[39m=\u001B[39;49muse_auth_token,\n\u001B[1;32m    420\u001B[0m         local_files_only\u001B[39m=\u001B[39;49mlocal_files_only,\n\u001B[1;32m    421\u001B[0m     )\n\u001B[1;32m    423\u001B[0m \u001B[39mexcept\u001B[39;00m RepositoryNotFoundError:\n\u001B[1;32m    424\u001B[0m     \u001B[39mraise\u001B[39;00m \u001B[39mEnvironmentError\u001B[39;00m(\n\u001B[1;32m    425\u001B[0m         \u001B[39mf\u001B[39m\u001B[39m\"\u001B[39m\u001B[39m{\u001B[39;00mpath_or_repo_id\u001B[39m}\u001B[39;00m\u001B[39m is not a local folder and is not a valid model identifier \u001B[39m\u001B[39m\"\u001B[39m\n\u001B[1;32m    426\u001B[0m         \u001B[39m\"\u001B[39m\u001B[39mlisted on \u001B[39m\u001B[39m'\u001B[39m\u001B[39mhttps://huggingface.co/models\u001B[39m\u001B[39m'\u001B[39m\u001B[39m\\n\u001B[39;00m\u001B[39mIf this is a private repository, make sure to \u001B[39m\u001B[39m\"\u001B[39m\n\u001B[1;32m    427\u001B[0m         \u001B[39m\"\u001B[39m\u001B[39mpass a token having permission to this repo with `use_auth_token` or log in with \u001B[39m\u001B[39m\"\u001B[39m\n\u001B[1;32m    428\u001B[0m         \u001B[39m\"\u001B[39m\u001B[39m`huggingface-cli login` and pass `use_auth_token=True`.\u001B[39m\u001B[39m\"\u001B[39m\n\u001B[1;32m    429\u001B[0m     )\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py:120\u001B[0m, in \u001B[0;36mvalidate_hf_hub_args.<locals>._inner_fn\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    117\u001B[0m \u001B[39mif\u001B[39;00m check_use_auth_token:\n\u001B[1;32m    118\u001B[0m     kwargs \u001B[39m=\u001B[39m smoothly_deprecate_use_auth_token(fn_name\u001B[39m=\u001B[39mfn\u001B[39m.\u001B[39m\u001B[39m__name__\u001B[39m, has_token\u001B[39m=\u001B[39mhas_token, kwargs\u001B[39m=\u001B[39mkwargs)\n\u001B[0;32m--> 120\u001B[0m \u001B[39mreturn\u001B[39;00m fn(\u001B[39m*\u001B[39;49margs, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs)\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/file_download.py:1134\u001B[0m, in \u001B[0;36mhf_hub_download\u001B[0;34m(repo_id, filename, subfolder, repo_type, revision, library_name, library_version, cache_dir, local_dir, local_dir_use_symlinks, user_agent, force_download, force_filename, proxies, etag_timeout, resume_download, token, local_files_only, legacy_cache_layout)\u001B[0m\n\u001B[1;32m   1132\u001B[0m \u001B[39mtry\u001B[39;00m:\n\u001B[1;32m   1133\u001B[0m     \u001B[39mtry\u001B[39;00m:\n\u001B[0;32m-> 1134\u001B[0m         metadata \u001B[39m=\u001B[39m get_hf_file_metadata(\n\u001B[1;32m   1135\u001B[0m             url\u001B[39m=\u001B[39;49murl,\n\u001B[1;32m   1136\u001B[0m             token\u001B[39m=\u001B[39;49mtoken,\n\u001B[1;32m   1137\u001B[0m             proxies\u001B[39m=\u001B[39;49mproxies,\n\u001B[1;32m   1138\u001B[0m             timeout\u001B[39m=\u001B[39;49metag_timeout,\n\u001B[1;32m   1139\u001B[0m         )\n\u001B[1;32m   1140\u001B[0m     \u001B[39mexcept\u001B[39;00m EntryNotFoundError \u001B[39mas\u001B[39;00m http_error:\n\u001B[1;32m   1141\u001B[0m         \u001B[39m# Cache the non-existence of the file and raise\u001B[39;00m\n\u001B[1;32m   1142\u001B[0m         commit_hash \u001B[39m=\u001B[39m http_error\u001B[39m.\u001B[39mresponse\u001B[39m.\u001B[39mheaders\u001B[39m.\u001B[39mget(HUGGINGFACE_HEADER_X_REPO_COMMIT)\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py:120\u001B[0m, in \u001B[0;36mvalidate_hf_hub_args.<locals>._inner_fn\u001B[0;34m(*args, **kwargs)\u001B[0m\n\u001B[1;32m    117\u001B[0m \u001B[39mif\u001B[39;00m check_use_auth_token:\n\u001B[1;32m    118\u001B[0m     kwargs \u001B[39m=\u001B[39m smoothly_deprecate_use_auth_token(fn_name\u001B[39m=\u001B[39mfn\u001B[39m.\u001B[39m\u001B[39m__name__\u001B[39m, has_token\u001B[39m=\u001B[39mhas_token, kwargs\u001B[39m=\u001B[39mkwargs)\n\u001B[0;32m--> 120\u001B[0m \u001B[39mreturn\u001B[39;00m fn(\u001B[39m*\u001B[39;49margs, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs)\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/file_download.py:1466\u001B[0m, in \u001B[0;36mget_hf_file_metadata\u001B[0;34m(url, token, proxies, timeout)\u001B[0m\n\u001B[1;32m   1463\u001B[0m headers \u001B[39m=\u001B[39m build_hf_headers(token\u001B[39m=\u001B[39mtoken)\n\u001B[1;32m   1465\u001B[0m \u001B[39m# Retrieve metadata\u001B[39;00m\n\u001B[0;32m-> 1466\u001B[0m r \u001B[39m=\u001B[39m _request_wrapper(\n\u001B[1;32m   1467\u001B[0m     method\u001B[39m=\u001B[39;49m\u001B[39m\"\u001B[39;49m\u001B[39mHEAD\u001B[39;49m\u001B[39m\"\u001B[39;49m,\n\u001B[1;32m   1468\u001B[0m     url\u001B[39m=\u001B[39;49murl,\n\u001B[1;32m   1469\u001B[0m     headers\u001B[39m=\u001B[39;49mheaders,\n\u001B[1;32m   1470\u001B[0m     allow_redirects\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m   1471\u001B[0m     follow_relative_redirects\u001B[39m=\u001B[39;49m\u001B[39mTrue\u001B[39;49;00m,\n\u001B[1;32m   1472\u001B[0m     proxies\u001B[39m=\u001B[39;49mproxies,\n\u001B[1;32m   1473\u001B[0m     timeout\u001B[39m=\u001B[39;49mtimeout,\n\u001B[1;32m   1474\u001B[0m )\n\u001B[1;32m   1475\u001B[0m hf_raise_for_status(r)\n\u001B[1;32m   1477\u001B[0m \u001B[39m# Return\u001B[39;00m\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/file_download.py:407\u001B[0m, in \u001B[0;36m_request_wrapper\u001B[0;34m(method, url, max_retries, base_wait_time, max_wait_time, timeout, follow_relative_redirects, **params)\u001B[0m\n\u001B[1;32m    405\u001B[0m \u001B[39m# 2. Force relative redirection\u001B[39;00m\n\u001B[1;32m    406\u001B[0m \u001B[39mif\u001B[39;00m follow_relative_redirects:\n\u001B[0;32m--> 407\u001B[0m     response \u001B[39m=\u001B[39m _request_wrapper(\n\u001B[1;32m    408\u001B[0m         method\u001B[39m=\u001B[39;49mmethod,\n\u001B[1;32m    409\u001B[0m         url\u001B[39m=\u001B[39;49murl,\n\u001B[1;32m    410\u001B[0m         max_retries\u001B[39m=\u001B[39;49mmax_retries,\n\u001B[1;32m    411\u001B[0m         base_wait_time\u001B[39m=\u001B[39;49mbase_wait_time,\n\u001B[1;32m    412\u001B[0m         max_wait_time\u001B[39m=\u001B[39;49mmax_wait_time,\n\u001B[1;32m    413\u001B[0m         timeout\u001B[39m=\u001B[39;49mtimeout,\n\u001B[1;32m    414\u001B[0m         follow_relative_redirects\u001B[39m=\u001B[39;49m\u001B[39mFalse\u001B[39;49;00m,\n\u001B[1;32m    415\u001B[0m         \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mparams,\n\u001B[1;32m    416\u001B[0m     )\n\u001B[1;32m    418\u001B[0m     \u001B[39m# If redirection, we redirect only relative paths.\u001B[39;00m\n\u001B[1;32m    419\u001B[0m     \u001B[39m# This is useful in case of a renamed repository.\u001B[39;00m\n\u001B[1;32m    420\u001B[0m     \u001B[39mif\u001B[39;00m \u001B[39m300\u001B[39m \u001B[39m<\u001B[39m\u001B[39m=\u001B[39m response\u001B[39m.\u001B[39mstatus_code \u001B[39m<\u001B[39m\u001B[39m=\u001B[39m \u001B[39m399\u001B[39m:\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/file_download.py:442\u001B[0m, in \u001B[0;36m_request_wrapper\u001B[0;34m(method, url, max_retries, base_wait_time, max_wait_time, timeout, follow_relative_redirects, **params)\u001B[0m\n\u001B[1;32m    439\u001B[0m     \u001B[39mreturn\u001B[39;00m response\n\u001B[1;32m    441\u001B[0m \u001B[39m# 3. Exponential backoff\u001B[39;00m\n\u001B[0;32m--> 442\u001B[0m \u001B[39mreturn\u001B[39;00m http_backoff(\n\u001B[1;32m    443\u001B[0m     method\u001B[39m=\u001B[39;49mmethod,\n\u001B[1;32m    444\u001B[0m     url\u001B[39m=\u001B[39;49murl,\n\u001B[1;32m    445\u001B[0m     max_retries\u001B[39m=\u001B[39;49mmax_retries,\n\u001B[1;32m    446\u001B[0m     base_wait_time\u001B[39m=\u001B[39;49mbase_wait_time,\n\u001B[1;32m    447\u001B[0m     max_wait_time\u001B[39m=\u001B[39;49mmax_wait_time,\n\u001B[1;32m    448\u001B[0m     retry_on_exceptions\u001B[39m=\u001B[39;49m(ConnectTimeout, ProxyError),\n\u001B[1;32m    449\u001B[0m     retry_on_status_codes\u001B[39m=\u001B[39;49m(),\n\u001B[1;32m    450\u001B[0m     timeout\u001B[39m=\u001B[39;49mtimeout,\n\u001B[1;32m    451\u001B[0m     \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mparams,\n\u001B[1;32m    452\u001B[0m )\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/huggingface_hub/utils/_http.py:129\u001B[0m, in \u001B[0;36mhttp_backoff\u001B[0;34m(method, url, max_retries, base_wait_time, max_wait_time, retry_on_exceptions, retry_on_status_codes, **kwargs)\u001B[0m\n\u001B[1;32m    126\u001B[0m     kwargs[\u001B[39m\"\u001B[39m\u001B[39mdata\u001B[39m\u001B[39m\"\u001B[39m]\u001B[39m.\u001B[39mseek(io_obj_initial_pos)\n\u001B[1;32m    128\u001B[0m \u001B[39m# Perform request and return if status_code is not in the retry list.\u001B[39;00m\n\u001B[0;32m--> 129\u001B[0m response \u001B[39m=\u001B[39m requests\u001B[39m.\u001B[39;49mrequest(method\u001B[39m=\u001B[39;49mmethod, url\u001B[39m=\u001B[39;49murl, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs)\n\u001B[1;32m    130\u001B[0m \u001B[39mif\u001B[39;00m response\u001B[39m.\u001B[39mstatus_code \u001B[39mnot\u001B[39;00m \u001B[39min\u001B[39;00m retry_on_status_codes:\n\u001B[1;32m    131\u001B[0m     \u001B[39mreturn\u001B[39;00m response\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/requests/api.py:59\u001B[0m, in \u001B[0;36mrequest\u001B[0;34m(method, url, **kwargs)\u001B[0m\n\u001B[1;32m     55\u001B[0m \u001B[39m# By using the 'with' statement we are sure the session is closed, thus we\u001B[39;00m\n\u001B[1;32m     56\u001B[0m \u001B[39m# avoid leaving sockets open which can trigger a ResourceWarning in some\u001B[39;00m\n\u001B[1;32m     57\u001B[0m \u001B[39m# cases, and look like a memory leak in others.\u001B[39;00m\n\u001B[1;32m     58\u001B[0m \u001B[39mwith\u001B[39;00m sessions\u001B[39m.\u001B[39mSession() \u001B[39mas\u001B[39;00m session:\n\u001B[0;32m---> 59\u001B[0m     \u001B[39mreturn\u001B[39;00m session\u001B[39m.\u001B[39;49mrequest(method\u001B[39m=\u001B[39;49mmethod, url\u001B[39m=\u001B[39;49murl, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs)\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/requests/sessions.py:587\u001B[0m, in \u001B[0;36mSession.request\u001B[0;34m(self, method, url, params, data, headers, cookies, files, auth, timeout, allow_redirects, proxies, hooks, stream, verify, cert, json)\u001B[0m\n\u001B[1;32m    582\u001B[0m send_kwargs \u001B[39m=\u001B[39m {\n\u001B[1;32m    583\u001B[0m     \u001B[39m\"\u001B[39m\u001B[39mtimeout\u001B[39m\u001B[39m\"\u001B[39m: timeout,\n\u001B[1;32m    584\u001B[0m     \u001B[39m\"\u001B[39m\u001B[39mallow_redirects\u001B[39m\u001B[39m\"\u001B[39m: allow_redirects,\n\u001B[1;32m    585\u001B[0m }\n\u001B[1;32m    586\u001B[0m send_kwargs\u001B[39m.\u001B[39mupdate(settings)\n\u001B[0;32m--> 587\u001B[0m resp \u001B[39m=\u001B[39m \u001B[39mself\u001B[39;49m\u001B[39m.\u001B[39;49msend(prep, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49msend_kwargs)\n\u001B[1;32m    589\u001B[0m \u001B[39mreturn\u001B[39;00m resp\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/requests/sessions.py:701\u001B[0m, in \u001B[0;36mSession.send\u001B[0;34m(self, request, **kwargs)\u001B[0m\n\u001B[1;32m    698\u001B[0m start \u001B[39m=\u001B[39m preferred_clock()\n\u001B[1;32m    700\u001B[0m \u001B[39m# Send the request\u001B[39;00m\n\u001B[0;32m--> 701\u001B[0m r \u001B[39m=\u001B[39m adapter\u001B[39m.\u001B[39;49msend(request, \u001B[39m*\u001B[39;49m\u001B[39m*\u001B[39;49mkwargs)\n\u001B[1;32m    703\u001B[0m \u001B[39m# Total elapsed time of the request (approximately)\u001B[39;00m\n\u001B[1;32m    704\u001B[0m elapsed \u001B[39m=\u001B[39m preferred_clock() \u001B[39m-\u001B[39m start\n",
      "File \u001B[0;32m~/anaconda3/envs/wm/lib/python3.10/site-packages/requests/adapters.py:563\u001B[0m, in \u001B[0;36mHTTPAdapter.send\u001B[0;34m(self, request, stream, timeout, verify, cert, proxies)\u001B[0m\n\u001B[1;32m    559\u001B[0m         \u001B[39mraise\u001B[39;00m ProxyError(e, request\u001B[39m=\u001B[39mrequest)\n\u001B[1;32m    561\u001B[0m     \u001B[39mif\u001B[39;00m \u001B[39misinstance\u001B[39m(e\u001B[39m.\u001B[39mreason, _SSLError):\n\u001B[1;32m    562\u001B[0m         \u001B[39m# This branch is for urllib3 v1.22 and later.\u001B[39;00m\n\u001B[0;32m--> 563\u001B[0m         \u001B[39mraise\u001B[39;00m SSLError(e, request\u001B[39m=\u001B[39mrequest)\n\u001B[1;32m    565\u001B[0m     \u001B[39mraise\u001B[39;00m \u001B[39mConnectionError\u001B[39;00m(e, request\u001B[39m=\u001B[39mrequest)\n\u001B[1;32m    567\u001B[0m \u001B[39mexcept\u001B[39;00m ClosedPoolError \u001B[39mas\u001B[39;00m e:\n",
      "\u001B[0;31mSSLError\u001B[0m: HTTPSConnectionPool(host='huggingface.co', port=443): Max retries exceeded with url: /gpt2/resolve/main/tokenizer_config.json (Caused by SSLError(SSLEOFError(8, 'EOF occurred in violation of protocol (_ssl.c:997)')))"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['CURL_CA_BUNDLE'] = ''\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "lm_tokenizer = AutoTokenizer.from_pretrained('gpt2')\n",
    "lm_model = AutoModelForCausalLM.from_pretrained('gpt2').to('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "model = model.to('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /home/wanglean/wm/c4-train.00000-of-00512_sliced/train/cache-bf7c139e2a7a4790.arrow\n"
     ]
    }
   ],
   "source": [
    "c4_sliced_and_filted = load_from_disk('./c4-train.00000-of-00512_sliced')\n",
    "c4_sliced_and_filted = c4_sliced_and_filted['train'].shuffle(seed=args.sample_seed).select(\n",
    "    range(100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 396,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 选择一个prompt\n",
    "\n",
    "sample_idx = 98\n",
    "input_text = c4_sliced_and_filted[sample_idx]['text']\n",
    "tokenized_input = tokenizer(input_text, return_tensors='pt').to(model.device)\n",
    "tokenized_input = truncate(tokenized_input, max_length=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "min_length_processor = MinLengthLogitsProcessor(min_length=1000,\n",
    "                                                eos_token_id=tokenizer.eos_token_id)\n",
    "repetition_processor = RepetitionPenaltyLogitsProcessor(penalty=1.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 253,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "import watermarking.watermark_processors.message_models.lm_message_model\n",
    "watermarking.watermark_processors.message_models.lm_message_model = reload(watermarking.watermark_processors.message_models.lm_message_model)\n",
    "import watermarking.watermark_processors.message_model_processor\n",
    "watermarking.watermark_processors.message_model_processor = reload(watermarking.watermark_processors.message_model_processor)\n",
    "\n",
    "\n",
    "from watermarking.watermark_processors.message_models.lm_message_model import LMMessageModel\n",
    "from watermarking.watermark_processors.message_model_processor import WmProcessorMessageModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 449,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "lm_message_model = LMMessageModel(tokenizer=tokenizer,lm_model=model,lm_tokenizer=tokenizer,\n",
    "    delta = 1.1, lm_prefix_len=10, lm_topk=-1, message_code_len = 20,random_permutation_num=50)\n",
    "wm_precessor_message_model = WmProcessorMessageModel(message_model=lm_message_model,tokenizer=tokenizer,\n",
    "    encode_ratio=5,max_confidence_lbd=0.5,strategy='max_confidence', message=[42,34,54,665,226,329])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 450,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 生成带watermark的文本\n",
    "\n",
    "start_length = tokenized_input['input_ids'].shape[-1]\n",
    "wm_precessor_message_model.start_length = start_length\n",
    "output_tokens = model.generate(**tokenized_input, max_new_tokens=200, num_beams=4,\n",
    "                               logits_processor=LogitsProcessorList(\n",
    "                                   [min_length_processor, repetition_processor,\n",
    "                                    wm_precessor_message_model]))\n",
    "output_text = tokenizer.decode(output_tokens[0][tokenized_input['input_ids'].shape[-1]:],\n",
    "                               skip_special_tokens=True) # 生成的文本\n",
    "prefix_and_output_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True) # 带prompt的文本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "log_probs = wm_precessor_message_model.decode(output_text)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "log_probs[0] # message"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 431,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 定义random watermark\n",
    "import watermarking.watermark_processors.random_processor\n",
    "watermarking.watermark_processors.random_processor = reload(watermarking.watermark_processors.random_processor)\n",
    "\n",
    "from watermarking.watermark_processors.random_processor import WmProcessorRandom\n",
    "random_processor = WmProcessorRandom(message=[42,34,5465,34], tokenizer=tokenizer,delta=1.5,message_code_len = 20,\n",
    "     top_k=100, encode_ratio = 5.\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 456,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 生成带random watermark的文本\n",
    "random_processor.start_length = tokenized_input['input_ids'].shape[-1]\n",
    "t_output_tokens = model.generate(**tokenized_input, max_new_tokens=200, num_beams=4,\n",
    "                               logits_processor=LogitsProcessorList(\n",
    "                                   [min_length_processor, repetition_processor,\n",
    "                                    random_processor]))\n",
    "t_output_text = tokenizer.decode(t_output_tokens[0][tokenized_input['input_ids'].shape[-1]:],\n",
    "                               skip_special_tokens=True) # 生成的文本\n",
    "t_prefix_and_output_text = tokenizer.decode(t_output_tokens[0], skip_special_tokens=True) # 带prompt的文本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 457,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:00<00:00, 1537.55it/s]\n"
     ]
    }
   ],
   "source": [
    "t_log_probs = random_processor.decode(t_output_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 458,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor(42, device='cuda:0'), tensor(34, device='cuda:0')]"
      ]
     },
     "execution_count": 458,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_log_probs[0] # message"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 430,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# 生成不带time watermark的文本\n",
    "\n",
    "no_wm_output_tokens = model.generate(**tokenized_input, max_new_tokens=200, num_beams=4,\n",
    "                                     logits_processor=LogitsProcessorList(\n",
    "                                       [min_length_processor, repetition_processor]))\n",
    "no_wm_output_text = tokenizer.decode(no_wm_output_tokens[0][tokenized_input['input_ids'].shape[-1]:],\n",
    "                                     skip_special_tokens=True)\n",
    "no_wm_prefix_and_output_text = tokenizer.decode(no_wm_output_tokens[0], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading opt-2.7b tokenizer from /home/wanglean/llm-ckpts/tokenizers/opt-2.7b ...\n",
      "finished loading\n",
      "loading opt-2.7b model from /home/wanglean/llm-ckpts/models/opt-2.7b ...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60101427088d4ec69570f86054c465a0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finished loading\n"
     ]
    }
   ],
   "source": [
    "oracle_tokenizer = load_local_model_or_tokenizer('facebook/opt-2.7b', 'tokenizer')\n",
    "oracle_model = load_local_model_or_tokenizer('facebook/opt-2.7b', 'model')\n",
    "oracle_model = oracle_model.to('cuda:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 454,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.2013373374938965 3.3245599269866943\n"
     ]
    }
   ],
   "source": [
    "from watermarking.experiments.watermark import compute_ppl_single\n",
    "\n",
    "# 计算带水印的ppl\n",
    "loss, ppl = compute_ppl_single(prefix_and_output_text=prefix_and_output_text,\n",
    "                               oracle_model_name=args.oracle_model_name,\n",
    "                               output_text=output_text,\n",
    "                               oracle_model=oracle_model, oracle_tokenizer=oracle_tokenizer)\n",
    "print(loss, ppl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 455,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.1480250358581543 3.1519618034362793\n"
     ]
    }
   ],
   "source": [
    "from watermarking.experiments.watermark import compute_ppl_single\n",
    "\n",
    "# 计算带time水印的ppl\n",
    "loss, ppl = compute_ppl_single(prefix_and_output_text=t_prefix_and_output_text,\n",
    "                               oracle_model_name=args.oracle_model_name,\n",
    "                               output_text=t_output_text,\n",
    "                               oracle_model=oracle_model, oracle_tokenizer=oracle_tokenizer)\n",
    "print(loss, ppl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 424,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9888984560966492 2.6882715225219727\n"
     ]
    }
   ],
   "source": [
    "# 计算不带水印的ppl\n",
    "loss, ppl = compute_ppl_single(prefix_and_output_text=no_wm_prefix_and_output_text,\n",
    "                               oracle_model_name=args.oracle_model_name,\n",
    "                               output_text=no_wm_output_text,\n",
    "                               oracle_model=oracle_model, oracle_tokenizer=oracle_tokenizer)\n",
    "print(loss, ppl)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.10 ('wm')",
   "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.10.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "69f52fabb15766d39c6bf90ba53c555c905cb082f5a671ecb5c4487727b3f015"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}