{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "80148df1-fd6b-4006-a55b-c40633f8a474",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'scripts.Constants'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[5], line 6\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mrandom\u001b[39;00m\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mcopy\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m deepcopy\n\u001b[1;32m----> 6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mscripts\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mConstants\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m EOS, UNK\n\u001b[0;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mscripts\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m save_text\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'scripts.Constants'"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer\n",
    "import random\n",
    "from copy import deepcopy\n",
    "# Constants\n",
    "EOS = 50256  # GPT-2's EOS token ID\n",
    "UNK = 50256  # Using EOS as UNK since GPT-2 doesn't have a specific UNK token\n",
    "\n",
    "# Utils\n",
    "def save_text(texts, path):\n",
    "    with open(path, 'w', encoding='utf-8') as f:\n",
    "        for text in texts:\n",
    "            f.write(text + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b0cdf70-c999-44a5-9fba-129b43540fc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Node:\n",
    "    def __init__(self, freq):\n",
    "        self.left = None\n",
    "        self.right = None\n",
    "        self.father = None\n",
    "        self.freq = freq\n",
    "    def isLeft(self):\n",
    "        return self.father.left == self\n",
    "\n",
    "def createNodes(freqs):\n",
    "    return [Node(freq) for freq in freqs]\n",
    "\n",
    "def createHuffmanTree(nodes):\n",
    "    queue = nodes[:]\n",
    "    while len(queue) > 1:\n",
    "        queue.sort(key=lambda item:item.freq)\n",
    "        node_left = queue.pop(0)\n",
    "        node_right = queue.pop(0)\n",
    "        node_father = Node(node_left.freq + node_right.freq)\n",
    "        node_father.left = node_left\n",
    "        node_father.right = node_right\n",
    "        node_left.father = node_father\n",
    "        node_right.father = node_father\n",
    "        queue.append(node_father)\n",
    "    queue[0].father = None\n",
    "    return queue[0]\n",
    "\n",
    "def huffmanEncoding(nodes, root):\n",
    "    codes = [''] * len(nodes)\n",
    "    for i in range(len(nodes)):\n",
    "        node_tmp = nodes[i]\n",
    "        while node_tmp != root:\n",
    "            if node_tmp.isLeft():\n",
    "                codes[i] = '0' + codes[i]\n",
    "            else:\n",
    "                codes[i] = '1' + codes[i]\n",
    "            node_tmp = node_tmp.father\n",
    "    return codes\n",
    "\n",
    "class Huffman():\n",
    "    def __init__(self, k=None, bits=None):\n",
    "        self.k = k\n",
    "        self.bits = bits\n",
    "        \n",
    "    def extract(self, pred_prob, word_index):\n",
    "        bit = None\n",
    "        prob, idx = torch.topk(pred_prob, k=self.k, dim=1)\n",
    "        word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]\n",
    "        nodes = createNodes([item[1] for item in word_prob])\n",
    "        root = createHuffmanTree(nodes)\n",
    "        codes = huffmanEncoding(nodes, root)\n",
    "        for i, w_p in enumerate(word_prob):\n",
    "            if w_p[0] == word_index:\n",
    "                bit = codes[i]\n",
    "                break\n",
    "        return bit\n",
    "    \n",
    "    def __call__(self, pred_prob):\n",
    "        prob, idx = torch.topk(pred_prob, k=self.k, dim=1)\n",
    "        word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]\n",
    "        nodes = createNodes([item[1] for item in word_prob])\n",
    "        root = createHuffmanTree(nodes)\n",
    "        codes = huffmanEncoding(nodes, root)\n",
    "        for i,code in enumerate(codes):\n",
    "            bit = self.bits[:len(code)]\n",
    "            if bit == code:\n",
    "                bit_word_index = word_prob[i][0]\n",
    "                self.bits = self.bits[len(code):]\n",
    "                break\n",
    "        if self.extract(pred_prob,bit_word_index)!=bit:\n",
    "            print(\"False\")\n",
    "        bit_word_index = torch.LongTensor([bit_word_index]).to(pred_prob.device)\n",
    "        return bit_word_index\n",
    "\n",
    "class FLC():\n",
    "    def __init__(self, k=None, bits=None):\n",
    "        self.k = k\n",
    "        self.bits = bits\n",
    "    \n",
    "    def extract(self, pred_prob, word_index):\n",
    "        bit = None\n",
    "        prob, idx = torch.topk(pred_prob, k=2**self.k, dim=1)\n",
    "        word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]\n",
    "        codes = [str(bin(i))[2:].zfill(self.k) for i in range(2**self.k)]\n",
    "        for i, w_p in enumerate(word_prob):\n",
    "            if w_p[0] == word_index:\n",
    "                bit = codes[i]\n",
    "                break\n",
    "        return bit\n",
    "    \n",
    "    def __call__(self, pred_prob):\n",
    "        prob, idx = torch.topk(pred_prob, k=2**self.k, dim=1)\n",
    "        word_prob = [[i.item(),j.item()] for i,j in zip(idx[0], prob[0])]\n",
    "        codes = [str(bin(i))[2:].zfill(self.k) for i in range(2**self.k)]\n",
    "        bit = self.bits[:self.k]\n",
    "        for i,code in enumerate(codes):\n",
    "            if bit==code:\n",
    "                bit_word_index = word_prob[i][0]\n",
    "                self.bits = self.bits[self.k:]\n",
    "                break\n",
    "        bit_word_index = torch.LongTensor([bit_word_index]).to(pred_prob.device)\n",
    "        return bit_word_index\n",
    "\n",
    "def setup_seed(seed):\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    random.seed(seed)\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "def generate_text_with_steganography(model, tokenizer, prompt, handle, max_length=30):\n",
    "    device = model.device\n",
    "    input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)\n",
    "    \n",
    "    # Generate steganographic text\n",
    "    stega_output = input_ids.clone()\n",
    "    normal_output = input_ids.clone()\n",
    "    \n",
    "    for _ in range(max_length):\n",
    "        if len(handle.bits) == 0:\n",
    "            break\n",
    "            \n",
    "        # Get model predictions\n",
    "        outputs = model(stega_output)\n",
    "        next_token_logits = outputs.logits[:, -1, :]\n",
    "        next_token_probs = F.softmax(next_token_logits, dim=-1)\n",
    "        \n",
    "        # Apply steganography\n",
    "        bit_word_index = handle(next_token_probs)\n",
    "        stega_output = torch.cat([stega_output, bit_word_index.unsqueeze(1)], dim=1)\n",
    "        \n",
    "        # Generate normal text\n",
    "        outputs = model(normal_output)\n",
    "        next_token_logits = outputs.logits[:, -1, :]\n",
    "        next_token_probs = F.softmax(next_token_logits, dim=-1)\n",
    "        next_token = torch.multinomial(next_token_probs, num_samples=1)\n",
    "        normal_output = torch.cat([normal_output, next_token], dim=1)\n",
    "        \n",
    "        if next_token[0].item() == tokenizer.eos_token_id:\n",
    "            break\n",
    "    \n",
    "    return stega_output, normal_output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a63e3b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def main():\n",
    "    setup_seed(2023)\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    # Load GPT-2 model and tokenizer\n",
    "    model_name = \"gpt2\"\n",
    "    tokenizer = GPT2Tokenizer.from_pretrained(model_name)\n",
    "    model = GPT2LMHeadModel.from_pretrained(model_name).to(device)\n",
    "    \n",
    "    # Generate random bits\n",
    "    bits = ''.join([str(random.choice(list(range(2)))) for i in range(3000)])\n",
    "    \n",
    "    # Initialize steganography handler\n",
    "    k = 2\n",
    "    handle = Huffman(k=2**k, bits=bits)\n",
    "    # handle = FLC(k=k, bits=bits)\n",
    "    \n",
    "    # Generate text with steganography\n",
    "    prompt = \"Once upon a time\"\n",
    "    stega_output, normal_output = generate_text_with_steganography(\n",
    "        model, tokenizer, prompt, handle\n",
    "    )\n",
    "    \n",
    "    # Decode and save results\n",
    "    stega_text = tokenizer.decode(stega_output[0], skip_special_tokens=True)\n",
    "    normal_text = tokenizer.decode(normal_output[0], skip_special_tokens=True)\n",
    "    \n",
    "    save_text([stega_text], './outputs-gpt2-stega.txt')\n",
    "    save_text([normal_text], './outputs-gpt2-normal.txt')\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}