adding test suite -- first commit

src/test/retool_genertion_with_cache_test.ipynb

@@ -0,0 +1,1585 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "ejiXlq27sck1",
+        "outputId": "d2c846e5-97da-4533-d23f-1cb876d67069"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Requirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (4.52.4)\n",
+            "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from transformers) (3.18.0)\n",
+            "Requirement already satisfied: huggingface-hub<1.0,>=0.30.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.32.4)\n",
+            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (2.0.2)\n",
+            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (24.2)\n",
+            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.11/dist-packages (from transformers) (6.0.2)\n",
+            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (2024.11.6)\n",
+            "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from transformers) (2.32.3)\n",
+            "Requirement already satisfied: tokenizers<0.22,>=0.21 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.21.1)\n",
+            "Requirement already satisfied: safetensors>=0.4.3 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.5.3)\n",
+            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers) (4.67.1)\n",
+            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.30.0->transformers) (2025.3.2)\n",
+            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.30.0->transformers) (4.14.0)\n",
+            "Requirement already satisfied: hf-xet<2.0.0,>=1.1.2 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.30.0->transformers) (1.1.2)\n",
+            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.4.2)\n",
+            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.10)\n",
+            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2.4.0)\n",
+            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2025.4.26)\n"
+          ]
+        }
+      ],
+      "source": [
+        "! pip install transformers"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "! pip install profiling-decorator"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "3Sa_Bpi1srA9",
+        "outputId": "6ad4ffd6-1058-4097-acb2-21978fe27ca0"
+      },
+      "execution_count": 2,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Collecting profiling-decorator\n",
+            "  Downloading profiling_decorator-0.0.6-py3-none-any.whl.metadata (6.2 kB)\n",
+            "Downloading profiling_decorator-0.0.6-py3-none-any.whl (9.2 kB)\n",
+            "Installing collected packages: profiling-decorator\n",
+            "Successfully installed profiling-decorator-0.0.6\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "9IRlvyF-J4Mp"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "eV3CXXy6J47P"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def test_updated_retool_implementation():\n",
+        "    # 1. Setup model, tokenizer, and device\n",
+        "    from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+        "    import torch\n",
+        "    import transformers\n",
+        "    import re\n",
+        "    from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache\n",
+        "\n",
+        "    # Use a model that fits in memory\n",
+        "    model_name = \"gpt2-medium\"\n",
+        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+        "\n",
+        "    # Ensure padding token is set\n",
+        "    if tokenizer.pad_token is None:\n",
+        "        tokenizer.pad_token = tokenizer.eos_token\n",
+        "\n",
+        "    # Check device\n",
+        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "    print(f\"Using device: {device}\")\n",
+        "\n",
+        "    # Load model to device\n",
+        "    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
+        "\n",
+        "    # 2. Add special tokens\n",
+        "    special_tokens = {\n",
+        "        'additional_special_tokens': ['<code>', '</code>', '<interpreter>', '</interpreter>']\n",
+        "    }\n",
+        "    tokenizer.add_special_tokens(special_tokens)\n",
+        "    model.resize_token_embeddings(len(tokenizer))\n",
+        "\n",
+        "    # Get token IDs\n",
+        "    code_start_id = tokenizer.convert_tokens_to_ids('<code>')\n",
+        "    code_end_id = tokenizer.convert_tokens_to_ids('</code>')\n",
+        "    interpreter_start_id = tokenizer.convert_tokens_to_ids('<interpreter>')\n",
+        "    interpreter_end_id = tokenizer.convert_tokens_to_ids('</interpreter>')\n",
+        "\n",
+        "    print(f\"EOS token ID: {tokenizer.eos_token_id}\")\n",
+        "    print(f\"Pad token ID: {tokenizer.pad_token_id}\")\n",
+        "    print(f\"Code tokens: {code_start_id}, {code_end_id}\")\n",
+        "    print(f\"Interpreter tokens: {interpreter_start_id}, {interpreter_end_id}\")\n",
+        "\n",
+        "    # 3. Create a test version of your ReToolTrainer with custom generation\n",
+        "    class TestReToolTrainer:\n",
+        "        def __init__(self, model, tokenizer, device):\n",
+        "            self.model = model\n",
+        "            self.processing_class = tokenizer\n",
+        "            self.device = device\n",
+        "            self.temperature = 0.7\n",
+        "            self.top_p = 0.9\n",
+        "            self.top_k = 50\n",
+        "\n",
+        "            # Ensure pad token is set\n",
+        "            if self.processing_class.pad_token is None:\n",
+        "                self.processing_class.pad_token = self.processing_class.eos_token\n",
+        "\n",
+        "        def _execute_code(self, code_block):\n",
+        "            \"\"\"Mock code execution\"\"\"\n",
+        "            print(f\"\\n==== EXECUTING CODE ====\")\n",
+        "            print(f\"{code_block}\")\n",
+        "            print(f\"========================\\n\")\n",
+        "            return \"0 1 1 2 3\"\n",
+        "\n",
+        "        def _custom_generate(self, input_ids, attention_mask=None, past_key_values=None, max_new_tokens=50, eos_token_ids=None):\n",
+        "            \"\"\"Custom generation function that avoids KV cache issues\"\"\"\n",
+        "            if attention_mask is None:\n",
+        "                attention_mask = torch.ones_like(input_ids)\n",
+        "\n",
+        "            if eos_token_ids is None:\n",
+        "                eos_token_ids = [self.processing_class.eos_token_id]\n",
+        "\n",
+        "            # Initialize\n",
+        "            current_ids = input_ids.clone()\n",
+        "            current_mask = attention_mask.clone()\n",
+        "            current_kv = past_key_values\n",
+        "\n",
+        "            # Generate tokens in batches for efficiency\n",
+        "            all_tokens = []\n",
+        "            batch_size = 10  # Process this many tokens at once\n",
+        "\n",
+        "            for start_idx in range(0, max_new_tokens, batch_size):\n",
+        "                # How many tokens to generate in this batch\n",
+        "                batch_tokens = min(batch_size, max_new_tokens - start_idx)\n",
+        "\n",
+        "                # Accumulate new tokens\n",
+        "                new_tokens = []\n",
+        "\n",
+        "                for _ in range(batch_tokens):\n",
+        "                    # Forward pass with proper cache handling\n",
+        "                    with torch.no_grad():\n",
+        "                        outputs = self.model(\n",
+        "                            input_ids=current_ids if current_kv is None else current_ids[:, -1:],\n",
+        "                            attention_mask=current_mask if current_kv is None else current_mask[:, -1:],\n",
+        "                            past_key_values=DynamicCache.from_legacy_cache(current_kv) if current_kv is not None else None,\n",
+        "                            use_cache=True\n",
+        "                        )\n",
+        "\n",
+        "                    # Sample next token\n",
+        "                    next_token_logits = outputs.logits[:, -1, :] / self.temperature\n",
+        "                    filtered_logits = self._filter_logits(next_token_logits)\n",
+        "                    probs = torch.nn.functional.softmax(filtered_logits, dim=-1)\n",
+        "                    next_token = torch.multinomial(probs, num_samples=1)\n",
+        "\n",
+        "                    # Add to accumulated tokens\n",
+        "                    token_id = next_token.item()\n",
+        "                    new_tokens.append(token_id)\n",
+        "\n",
+        "                    # Update for next iteration\n",
+        "                    current_ids = torch.cat([current_ids, next_token], dim=1)\n",
+        "                    token_mask = torch.ones((1, 1), device=current_mask.device, dtype=current_mask.dtype)\n",
+        "                    current_mask = torch.cat([current_mask, token_mask], dim=1)\n",
+        "                    current_kv = outputs.past_key_values\n",
+        "\n",
+        "                    # Check for stop tokens - include both EOS and code_end\n",
+        "                    if token_id in eos_token_ids:\n",
+        "                        break\n",
+        "\n",
+        "                # Add batch tokens to overall result\n",
+        "                all_tokens.extend(new_tokens)\n",
+        "\n",
+        "                # Check if we hit a stop token\n",
+        "                if len(new_tokens) < batch_tokens:\n",
+        "                    break\n",
+        "\n",
+        "            # Convert to tensor\n",
+        "            result = torch.tensor([all_tokens], device=input_ids.device)\n",
+        "            return result, current_kv\n",
+        "\n",
+        "        def _filter_logits(self, logits):\n",
+        "            \"\"\"Apply top-k and top-p filtering\"\"\"\n",
+        "            if self.top_k > 0:\n",
+        "                top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)\n",
+        "                logits[0, :] = torch.full_like(logits[0, :], float('-inf'))\n",
+        "                logits[0, top_k_indices[0]] = top_k_logits[0]\n",
+        "\n",
+        "            if self.top_p < 1.0:\n",
+        "                sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)\n",
+        "                cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)\n",
+        "\n",
+        "                # Remove tokens with cumulative probability above threshold\n",
+        "                sorted_indices_to_remove = cumulative_probs > self.top_p\n",
+        "                # Shift the indices to the right to keep the first token above threshold\n",
+        "                sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()\n",
+        "                sorted_indices_to_remove[:, 0] = 0\n",
+        "\n",
+        "                # Scatter sorted tensors to original indexing\n",
+        "                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n",
+        "                logits[indices_to_remove] = float('-inf')\n",
+        "\n",
+        "            return logits\n",
+        "\n",
+        "        def _retool_generate_with_interpreter(self, prompt_ids_batch, attention_mask_batch, eos_id, interpreter_id, code_id, max_turns=10):\n",
+        "            \"\"\"Your updated implementation with custom generation\"\"\"\n",
+        "            batch_size = prompt_ids_batch.size(0)\n",
+        "            batch_completion = []\n",
+        "            batch_interpreter_positions = []\n",
+        "\n",
+        "            for i in range(batch_size):\n",
+        "                print(f\"Processing batch item {i+1}/{batch_size}\")\n",
+        "\n",
+        "                # Initialize\n",
+        "                current_input_id = prompt_ids_batch[i:i+1]\n",
+        "                current_attention_mask = attention_mask_batch[i:i+1]\n",
+        "                current_kv = None\n",
+        "\n",
+        "                # Track the completion part (no prompt)\n",
+        "                cumulative_completion_ids = torch.empty((1, 0), dtype=torch.long, device=prompt_ids_batch.device)\n",
+        "                interpreter_positions = []\n",
+        "\n",
+        "                for turn_idx in range(max_turns):\n",
+        "                    # Check if input is empty\n",
+        "                    if current_input_id.size(1) == 0:\n",
+        "                        print(f\"Turn {turn_idx + 1}: Input is empty, breaking loop\")\n",
+        "                        break\n",
+        "\n",
+        "                    print(f\"\\n--- Turn {turn_idx + 1} ---\")\n",
+        "                    print(f\"Current input: {self.processing_class.decode(current_input_id[0])}\")\n",
+        "                    print(f\"KV cache present: {current_kv is not None}\")\n",
+        "\n",
+        "                    # Generate with custom function\n",
+        "                    newly_generated_tokens, current_kv = self._custom_generate(\n",
+        "                        input_ids=current_input_id,\n",
+        "                        attention_mask=current_attention_mask,\n",
+        "                        past_key_values=current_kv,\n",
+        "                        max_new_tokens=30,\n",
+        "                        eos_token_ids=[eos_id, code_id[1]]\n",
+        "                    )\n",
+        "\n",
+        "                    # Display generated text\n",
+        "                    print(f\"Generated: {self.processing_class.decode(newly_generated_tokens[0])}\")\n",
+        "\n",
+        "                    # Add to cumulative completion\n",
+        "                    cumulative_completion_ids = torch.cat([cumulative_completion_ids, newly_generated_tokens], dim=1)\n",
+        "\n",
+        "                    # Check last token\n",
+        "                    last_token_id = newly_generated_tokens[0, -1].item() if newly_generated_tokens.size(1) > 0 else None\n",
+        "                    print(f\"Last token ID: {last_token_id}\")\n",
+        "\n",
+        "                    # Check for end conditions\n",
+        "                    if last_token_id == eos_id:\n",
+        "                        print(\"Found EOS token, ending generation\")\n",
+        "                        break\n",
+        "\n",
+        "                    # Check for code end token\n",
+        "                    if last_token_id == code_id[1]:\n",
+        "                        print(\"Found </code> token, executing code\")\n",
+        "\n",
+        "                        # Extract code from the full text\n",
+        "                        full_text = self.processing_class.decode(\n",
+        "                            torch.cat([prompt_ids_batch[i], cumulative_completion_ids[0]], dim=0)\n",
+        "                        )\n",
+        "                        code_match = re.search(r'<code>(.*?)</code>', full_text, re.DOTALL)\n",
+        "\n",
+        "                        if code_match:\n",
+        "                            code_block = code_match.group(1).strip()\n",
+        "\n",
+        "                            # Execute code\n",
+        "                            interpreter_text = self._execute_code(code_block)\n",
+        "\n",
+        "                            # Format and add interpreter output\n",
+        "                            formatted_feedback = f\"{self.processing_class.decode(interpreter_id[0])}{interpreter_text}{self.processing_class.decode(interpreter_id[1])}\"\n",
+        "                            interpreter_ids = self.processing_class(\n",
+        "                                formatted_feedback,\n",
+        "                                return_tensors=\"pt\",\n",
+        "                                add_special_tokens=False\n",
+        "                            ).input_ids.to(prompt_ids_batch.device)\n",
+        "\n",
+        "                            # Record positions\n",
+        "                            interpreter_start_idx = cumulative_completion_ids.size(1)\n",
+        "                            cumulative_completion_ids = torch.cat([cumulative_completion_ids, interpreter_ids], dim=1)\n",
+        "                            interpreter_end_idx = cumulative_completion_ids.size(1) - 1\n",
+        "                            interpreter_positions.append((interpreter_start_idx, interpreter_end_idx))\n",
+        "\n",
+        "                            print(f\"Added interpreter output: {formatted_feedback}\")\n",
+        "\n",
+        "                            # Set up for next turn\n",
+        "                            current_input_id = interpreter_ids\n",
+        "                            current_attention_mask = torch.ones_like(current_input_id)\n",
+        "                            # Keep current_kv from previous generation\n",
+        "                        else:\n",
+        "                            print(\"No code block found despite </code> token\")\n",
+        "                            break\n",
+        "                    else:\n",
+        "                        # Continue with the newly generated tokens\n",
+        "                        current_input_id = newly_generated_tokens\n",
+        "                        current_attention_mask = torch.ones_like(current_input_id)\n",
+        "\n",
+        "                # Add to batch results\n",
+        "                batch_completion.append(cumulative_completion_ids.squeeze(0))\n",
+        "                batch_interpreter_positions.append(interpreter_positions)\n",
+        "\n",
+        "            # Pad sequences\n",
+        "            if len(batch_completion) > 0:\n",
+        "                # Ensure padding_value is a valid integer\n",
+        "                padding_value = self.processing_class.pad_token_id\n",
+        "                if padding_value is None:\n",
+        "                    padding_value = 0  # Use 0 as a default if pad_token_id is None\n",
+        "\n",
+        "                padded_sequences = torch.nn.utils.rnn.pad_sequence(\n",
+        "                    batch_completion,\n",
+        "                    batch_first=True,\n",
+        "                    padding_value=padding_value\n",
+        "                )\n",
+        "            else:\n",
+        "                padded_sequences = torch.empty((0, 0), dtype=torch.long, device=prompt_ids_batch.device)\n",
+        "\n",
+        "            return padded_sequences, batch_interpreter_positions\n",
+        "\n",
+        "    # 4. Create test instance\n",
+        "    tester = TestReToolTrainer(model, tokenizer, device)\n",
+        "\n",
+        "    # 5. Create a test prompt with a complete code block\n",
+        "    prompt = \"\"\"Let me solve this problem with code:\n",
+        "\n",
+        "<code>\n",
+        "def fibonacci(n):\n",
+        "    a, b = 0, 1\n",
+        "    result = []\n",
+        "    for _ in range(n):\n",
+        "        result.append(a)\n",
+        "        a, b = b, a + b\n",
+        "    return result\n",
+        "\n",
+        "print(fibonacci(5))\n",
+        "</code>\"\"\"\n",
+        "\n",
+        "    # 6. Run the test\n",
+        "    try:\n",
+        "        print(\"\\n=== Testing Updated ReTool Implementation ===\\n\")\n",
+        "\n",
+        "        # Encode the prompt\n",
+        "        prompt_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
+        "        attention_mask = torch.ones_like(prompt_ids)\n",
+        "\n",
+        "        # Run the generation\n",
+        "        completions, positions = tester._retool_generate_with_interpreter(\n",
+        "            prompt_ids_batch=prompt_ids,\n",
+        "            attention_mask_batch=attention_mask,\n",
+        "            eos_id=tokenizer.eos_token_id,\n",
+        "            interpreter_id=[interpreter_start_id, interpreter_end_id],\n",
+        "            code_id=[code_start_id, code_end_id],\n",
+        "            max_turns=3\n",
+        "        )\n",
+        "\n",
+        "        # Display results\n",
+        "        print(\"\\n=== Final Results ===\\n\")\n",
+        "        print(\"Generated completion:\")\n",
+        "        print(tokenizer.decode(completions[0]))\n",
+        "\n",
+        "        print(\"\\nFull text:\")\n",
+        "        print(tokenizer.decode(torch.cat([prompt_ids[0], completions[0]])))\n",
+        "\n",
+        "        print(\"\\nInterpreter positions:\", positions)\n",
+        "\n",
+        "    except Exception as e:\n",
+        "        import traceback\n",
+        "        print(f\"Error during testing: {e}\")\n",
+        "        traceback.print_exc()\n",
+        "\n",
+        "# Run the test\n",
+        "test_updated_retool_implementation()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "4_E6Eo7EHC_8",
+        "outputId": "35b195d9-b0ff-4ddf-c216-fba1c83f40e2"
+      },
+      "execution_count": 9,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Using device: cpu\n",
+            "EOS token ID: 50256\n",
+            "Pad token ID: 50256\n",
+            "Code tokens: 50257, 50258\n",
+            "Interpreter tokens: 50259, 50260\n",
+            "\n",
+            "=== Testing Updated ReTool Implementation ===\n",
+            "\n",
+            "Processing batch item 1/1\n",
+            "\n",
+            "--- Turn 1 ---\n",
+            "Current input: Let me solve this problem with code:\n",
+            "\n",
+            "<code>\n",
+            "def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "</code>\n",
+            "KV cache present: False\n",
+            "Generated: \n",
+            "\n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = []\n",
+            "\n",
+            "Last token ID: 198\n",
+            "\n",
+            "--- Turn 2 ---\n",
+            "Current input: \n",
+            "\n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = []\n",
+            "\n",
+            "KV cache present: True\n",
+            "Generated: \n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "    return result\n",
+            "\n",
+            "print(fibon\n",
+            "Last token ID: 261\n",
+            "\n",
+            "--- Turn 3 ---\n",
+            "Current input: \n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "    return result\n",
+            "\n",
+            "print(fibon\n",
+            "KV cache present: True\n",
+            "Generated: acci(5))\n",
+            "\n",
+            "So the first two methods are all the same, the last one is a little different, and the second one is the\n",
+            "Last token ID: 262\n",
+            "\n",
+            "=== Final Results ===\n",
+            "\n",
+            "Generated completion:\n",
+            "\n",
+            "\n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = []\n",
+            "\n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "\n",
+            "So the first two methods are all the same, the last one is a little different, and the second one is the\n",
+            "\n",
+            "Full text:\n",
+            "Let me solve this problem with code:\n",
+            "\n",
+            "<code>\n",
+            "def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "</code>\n",
+            "\n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = []\n",
+            "\n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "\n",
+            "So the first two methods are all the same, the last one is a little different, and the second one is the\n",
+            "\n",
+            "Interpreter positions: [[]]\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "Z0EHHkP3J7Ox"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def test_retool_core_functionality():\n",
+        "    # 1. Create minimal model and tokenizer\n",
+        "    from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+        "    import torch\n",
+        "    import transformers\n",
+        "    import re\n",
+        "\n",
+        "    # Use a small model for testing\n",
+        "    model_name = \"gpt2-medium\"\n",
+        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+        "\n",
+        "    # Check if CUDA is available\n",
+        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "    print(f\"Using device: {device}\")\n",
+        "\n",
+        "    # Load model directly to the selected device\n",
+        "    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
+        "\n",
+        "    # 2. Add special tokens to the tokenizer\n",
+        "    special_tokens = {\n",
+        "        'additional_special_tokens': ['<code>', '</code>', '<interpreter>', '</interpreter>']\n",
+        "    }\n",
+        "    tokenizer.add_special_tokens(special_tokens)\n",
+        "    model.resize_token_embeddings(len(tokenizer))\n",
+        "\n",
+        "    # Get token IDs for special tokens\n",
+        "    code_start_id = tokenizer.convert_tokens_to_ids('<code>')\n",
+        "    code_end_id = tokenizer.convert_tokens_to_ids('</code>')\n",
+        "    interpreter_start_id = tokenizer.convert_tokens_to_ids('<interpreter>')\n",
+        "    interpreter_end_id = tokenizer.convert_tokens_to_ids('</interpreter>')\n",
+        "\n",
+        "    print(f\"Code tokens: {code_start_id}, {code_end_id}\")\n",
+        "    print(f\"Interpreter tokens: {interpreter_start_id}, {interpreter_end_id}\")\n",
+        "\n",
+        "    # 3. Create a simplified implementation of _retool_generate_with_interpreter\n",
+        "    def simplified_generate_with_interpreter(model, tokenizer, prompt_text, device):\n",
+        "        \"\"\"Simplified version focusing just on the core functionality\"\"\"\n",
+        "        # Step 1: Tokenize the prompt\n",
+        "        prompt_ids = tokenizer.encode(prompt_text, return_tensors=\"pt\").to(device)\n",
+        "\n",
+        "        # Initialize tracking variables\n",
+        "        cumulative_completion_ids = torch.empty((1, 0), dtype=torch.long, device=device)\n",
+        "        interpreter_positions = []\n",
+        "\n",
+        "        # Step 2: Extract a code block and execute it\n",
+        "        full_text = prompt_text\n",
+        "        code_match = re.search(r'<code>(.*?)</code>', full_text, re.DOTALL)\n",
+        "\n",
+        "        if code_match:\n",
+        "            code_block = code_match.group(1).strip()\n",
+        "            print(f\"Found code block: {code_block}\")\n",
+        "\n",
+        "            # Mock code execution\n",
+        "            interpreter_output = \"0 1 1 2 3\"\n",
+        "            print(f\"Code execution result: {interpreter_output}\")\n",
+        "\n",
+        "            # Format interpreter feedback\n",
+        "            interpreter_text = f\"<interpreter>{interpreter_output}</interpreter>\"\n",
+        "            interpreter_ids = tokenizer.encode(\n",
+        "                interpreter_text,\n",
+        "                return_tensors=\"pt\",\n",
+        "                add_special_tokens=False\n",
+        "            ).to(device)\n",
+        "\n",
+        "            # Step 3: Generate a continuation after the interpreter output\n",
+        "            # First, create a sequence with prompt + interpreter output\n",
+        "            combined_input = prompt_text + interpreter_text\n",
+        "            combined_ids = tokenizer.encode(combined_input, return_tensors=\"pt\").to(device)\n",
+        "\n",
+        "            # Generate continuation\n",
+        "            with torch.no_grad():\n",
+        "                continuation_outputs = model.generate(\n",
+        "                    input_ids=combined_ids,\n",
+        "                    max_new_tokens=50,\n",
+        "                    do_sample=True,\n",
+        "                    temperature=0.7,\n",
+        "                    top_p=0.9,\n",
+        "                    pad_token_id=tokenizer.pad_token_id,\n",
+        "                    eos_token_id=tokenizer.eos_token_id,\n",
+        "                    return_dict_in_generate=True,\n",
+        "                    cache_implementation= 'offloaded',\n",
+        "                )\n",
+        "\n",
+        "            # Extract only the newly generated tokens\n",
+        "            continuation_tokens = continuation_outputs.sequences[:, combined_ids.size(1):]\n",
+        "\n",
+        "            # Combine everything for the final result\n",
+        "            # The completion consists of: interpreter output + continuation\n",
+        "            cumulative_completion_ids = torch.cat([interpreter_ids, continuation_tokens], dim=1)\n",
+        "\n",
+        "            # Record the interpreter position\n",
+        "            interpreter_positions.append((0, interpreter_ids.size(1) - 1))\n",
+        "\n",
+        "            print(f\"Generated continuation: {tokenizer.decode(continuation_tokens[0])}\")\n",
+        "        else:\n",
+        "            print(\"No code block found in the prompt.\")\n",
+        "\n",
+        "        return cumulative_completion_ids, interpreter_positions\n",
+        "\n",
+        "    # 4. Test with a prompt that has a complete code block\n",
+        "    prompt = \"\"\"Let's calculate Fibonacci numbers in Python:\n",
+        "\n",
+        "<code>\n",
+        "def fibonacci(n):\n",
+        "    a, b = 0, 1\n",
+        "    result = []\n",
+        "    for _ in range(n):\n",
+        "        result.append(a)\n",
+        "        a, b = b, a + b\n",
+        "    return result\n",
+        "\n",
+        "print(fibonacci(5))\n",
+        "</code>\"\"\"\n",
+        "\n",
+        "    # 5. Run our simplified test\n",
+        "    try:\n",
+        "        print(\"\\n--- Testing Core Functionality ---\")\n",
+        "        completion, positions = simplified_generate_with_interpreter(model, tokenizer, prompt, device)\n",
+        "\n",
+        "        print(\"\\n--- Final Results ---\")\n",
+        "        print(\"Completion:\")\n",
+        "        print(tokenizer.decode(completion[0]))\n",
+        "        print(\"\\nInterpreter positions:\", positions)\n",
+        "\n",
+        "        # 6. Now also test ReToolTrainer to verify core code execution functionality\n",
+        "        print(\"\\n--- Testing ReToolTrainer with Direct Injection ---\")\n",
+        "\n",
+        "        # Setup trainer\n",
+        "        trainer = ReToolTrainer(\n",
+        "            model=model,\n",
+        "            processing_class=tokenizer,\n",
+        "            args=transformers.TrainingArguments(\n",
+        "                output_dir=\"./test_output\",\n",
+        "                per_device_train_batch_size=1,\n",
+        "            ),\n",
+        "            train_dataset=None,\n",
+        "            eval_dataset=None,\n",
+        "            max_turns=3,\n",
+        "            interpreter_id=[interpreter_start_id, interpreter_end_id],\n",
+        "            code_id=[code_start_id, code_end_id],\n",
+        "            eos_id=tokenizer.eos_token_id\n",
+        "        )\n",
+        "\n",
+        "        # Override the _execute_code method\n",
+        "        def mock_execute_code(self, code_block):\n",
+        "            print(f\"Mock executing code: {code_block}\")\n",
+        "            return \"0 1 1 2 3\"\n",
+        "\n",
+        "        original_execute_code = trainer._execute_code\n",
+        "        trainer._execute_code = mock_execute_code.__get__(trainer, ReToolTrainer)\n",
+        "\n",
+        "        # Create a sequence that has a prompt and ends with </code>\n",
+        "        # This is to simulate that the model has generated a complete code block\n",
+        "        prompt_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
+        "        attention_mask = torch.ones_like(prompt_ids)\n",
+        "\n",
+        "        # Directly inject a simulated generation with a code block\n",
+        "        # Create a custom testing function\n",
+        "        def test_execute_code_and_continue(self, prompt_ids, attention_mask):\n",
+        "            \"\"\"Test just the code execution and continuation part\"\"\"\n",
+        "            print(\"Testing code execution and continuation...\")\n",
+        "            device = next(self.model.parameters()).device\n",
+        "            prompt_ids = prompt_ids.to(device)\n",
+        "            attention_mask = attention_mask.to(device)\n",
+        "\n",
+        "            # Extract code from the prompt\n",
+        "            full_text = self.processing_class.decode(prompt_ids[0])\n",
+        "            code_match = re.search(r'<code>(.*?)</code>', full_text, re.DOTALL)\n",
+        "\n",
+        "            if not code_match:\n",
+        "                print(\"No code block found in the prompt!\")\n",
+        "                return None, []\n",
+        "\n",
+        "            code_block = code_match.group(1).strip()\n",
+        "            print(f\"Executing code block: {code_block}\")\n",
+        "\n",
+        "            # Execute the code\n",
+        "            interpreter_text = self._execute_code(code_block)\n",
+        "\n",
+        "            # Format and tokenize the interpreter output\n",
+        "            formatted_feedback = f\"{self.processing_class.decode(self.interpreter_id[0])}{interpreter_text}{self.processing_class.decode(self.interpreter_id[1])}\"\n",
+        "            interpreter_ids = self.processing_class(\n",
+        "                formatted_feedback,\n",
+        "                return_tensors=\"pt\",\n",
+        "                add_special_tokens=False\n",
+        "            ).input_ids.to(device)\n",
+        "\n",
+        "            # Record position (relative to completion only)\n",
+        "            interpreter_positions = [(0, interpreter_ids.size(1) - 1)]\n",
+        "\n",
+        "            # Combine prompt with interpreter output for continuation\n",
+        "            combined_ids = torch.cat([prompt_ids, interpreter_ids], dim=1)\n",
+        "            combined_mask = torch.ones_like(combined_ids)\n",
+        "\n",
+        "            # Generate continuation\n",
+        "            continuation_outputs = self.model.generate(\n",
+        "                input_ids=combined_ids,\n",
+        "                attention_mask=combined_mask,\n",
+        "                max_new_tokens=50,\n",
+        "                do_sample=True,\n",
+        "                temperature=0.7,\n",
+        "                top_p=0.9,\n",
+        "                pad_token_id=self.processing_class.pad_token_id,\n",
+        "                eos_token_id=self.eos_id,\n",
+        "                return_dict_in_generate=True,\n",
+        "                cache_implementation= 'offloaded',\n",
+        "            )\n",
+        "\n",
+        "            # Extract only the newly generated continuation\n",
+        "            continuation_tokens = continuation_outputs.sequences[:, combined_ids.size(1):]\n",
+        "\n",
+        "            # Full completion is: interpreter output + continuation\n",
+        "            completion = torch.cat([interpreter_ids, continuation_tokens], dim=1)\n",
+        "\n",
+        "            return completion, interpreter_positions\n",
+        "\n",
+        "        # Add the test method to the trainer\n",
+        "        trainer.test_execute_code_and_continue = test_execute_code_and_continue.__get__(trainer, ReToolTrainer)\n",
+        "\n",
+        "        # Run the test\n",
+        "        completion, positions = trainer.test_execute_code_and_continue(prompt_ids, attention_mask)\n",
+        "\n",
+        "        print(\"\\n--- Trainer Test Results ---\")\n",
+        "        if completion is not None:\n",
+        "            print(\"Completion:\")\n",
+        "            print(tokenizer.decode(completion[0]))\n",
+        "            print(\"\\nInterpreter positions:\", positions)\n",
+        "\n",
+        "    except Exception as e:\n",
+        "        import traceback\n",
+        "        print(f\"Error during testing: {e}\")\n",
+        "        traceback.print_exc()\n",
+        "    finally:\n",
+        "        # Restore original method if needed\n",
+        "        if 'trainer' in locals() and 'original_execute_code' in locals():\n",
+        "            trainer._execute_code = original_execute_code\n",
+        "\n",
+        "# Run the test\n",
+        "test_retool_core_functionality()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "MumcLxASaBkj",
+        "outputId": "21dd5d29-f402-4837-be86-08cee4b9d7a2"
+      },
+      "execution_count": 25,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Using device: cuda\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
+            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Code tokens: 50257, 50258\n",
+            "Interpreter tokens: 50259, 50260\n",
+            "\n",
+            "--- Testing Core Functionality ---\n",
+            "Found code block: def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "Code execution result: 0 1 1 2 3\n",
+            "Generated continuation: 0 1 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48\n",
+            "\n",
+            "--- Final Results ---\n",
+            "Completion:\n",
+            "<interpreter>0 1 1 2 3</interpreter>0 1 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48\n",
+            "\n",
+            "Interpreter positions: [(0, 6)]\n",
+            "\n",
+            "--- Testing ReToolTrainer with Direct Injection ---\n",
+            "Testing code execution and continuation...\n",
+            "Executing code block: def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "Mock executing code: def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "/tmp/ipython-input-20-2039368761.py:57: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `ReToolTrainer.__init__`. Use `processing_class` instead.\n",
+            "  super().__init__(\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\n",
+            "--- Trainer Test Results ---\n",
+            "Completion:\n",
+            "<interpreter>0 1 1 2 3</interpreter>1 1 1 2 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n",
+            "\n",
+            "Interpreter positions: [(0, 6)]\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "2KBJZTXOaCA3"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache\n",
+        "\n",
+        "def test_direct_kv_cache_usage():\n",
+        "    # 1. Setup model, tokenizer, and device\n",
+        "    from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+        "    import torch\n",
+        "\n",
+        "    # Use a model that fits in memory\n",
+        "    model_name = \"gpt2-medium\"\n",
+        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+        "\n",
+        "    # Check device\n",
+        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "    print(f\"Using device: {device}\")\n",
+        "\n",
+        "    # Load model to device\n",
+        "    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
+        "\n",
+        "    # 2. Manual token-by-token generation with KV caching\n",
+        "    def generate_with_manual_kv_cache(input_ids, num_tokens=20):\n",
+        "        \"\"\"Generate tokens one by one with manual KV cache management\"\"\"\n",
+        "        current_ids = input_ids.clone()\n",
+        "        past_key_values = None\n",
+        "\n",
+        "        generated_tokens = []\n",
+        "\n",
+        "        for _ in range(num_tokens):\n",
+        "            # Forward pass with past_key_values\n",
+        "            with torch.no_grad():\n",
+        "                outputs = model(\n",
+        "                    input_ids=current_ids if past_key_values is None else current_ids[:, -1:],\n",
+        "                    #past_key_values=past_key_values,\n",
+        "                    past_key_values= DynamicCache.from_legacy_cache(past_key_values),\n",
+        "                    use_cache=True\n",
+        "                )\n",
+        "\n",
+        "            # Get logits for the next token (last position)\n",
+        "            next_token_logits = outputs.logits[:, -1, :]\n",
+        "\n",
+        "            # Sample from the distribution\n",
+        "            probs = torch.nn.functional.softmax(next_token_logits / 0.7, dim=-1)\n",
+        "            next_token = torch.multinomial(probs, num_samples=1)\n",
+        "\n",
+        "            # Add to generated tokens\n",
+        "            generated_tokens.append(next_token.item())\n",
+        "\n",
+        "            # Update current_ids for next iteration\n",
+        "            current_ids = torch.cat([current_ids, next_token], dim=1)\n",
+        "\n",
+        "            # Update past_key_values\n",
+        "            past_key_values = outputs.past_key_values\n",
+        "            #print('after generation, past_key_values ', past_key_values)\n",
+        "\n",
+        "        return generated_tokens, past_key_values\n",
+        "\n",
+        "    # 3. Run multi-turn generation with manual KV cache\n",
+        "    def run_manual_multi_turn_generation():\n",
+        "        # Start with a prompt\n",
+        "        prompt = \"Once upon a time, in a magical forest, there lived a\"\n",
+        "\n",
+        "        # Tokenize the prompt\n",
+        "        prompt_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
+        "\n",
+        "        # Initialize tracking\n",
+        "        full_text = prompt\n",
+        "        current_ids = prompt_ids\n",
+        "        past_kv = None\n",
+        "\n",
+        "        # Generate in multiple turns\n",
+        "        for turn_idx in range(3):  # 3 turns\n",
+        "            print(f\"\\n==== Turn {turn_idx + 1} ====\")\n",
+        "            print(f\"Current input: {tokenizer.decode(current_ids[0])}\")\n",
+        "            print(f\"KV cache present: {past_kv is not None}\")\n",
+        "\n",
+        "            # Pause for inspection\n",
+        "            input(\"Press Enter to generate next part...\")\n",
+        "\n",
+        "            # Generate new tokens manually\n",
+        "            new_token_ids, past_kv = generate_with_manual_kv_cache(current_ids, num_tokens=20)\n",
+        "\n",
+        "            # Decode and display new tokens\n",
+        "            new_text = tokenizer.decode(new_token_ids)\n",
+        "            print(f\"Generated: {new_text}\")\n",
+        "\n",
+        "            # Accumulate\n",
+        "            full_text += new_text\n",
+        "\n",
+        "            # Now inject a custom continuation\n",
+        "            custom_text = \" Suddenly, a rainbow appeared in the sky!\"\n",
+        "            custom_ids = tokenizer.encode(custom_text, return_tensors=\"pt\").to(device)\n",
+        "\n",
+        "            print(f\"\\n==== Injecting custom text: {custom_text} ====\")\n",
+        "\n",
+        "            # Update tracking\n",
+        "            full_text += custom_text\n",
+        "\n",
+        "            # Prepare for next turn - start with the custom text\n",
+        "            current_ids = custom_ids\n",
+        "            # Keep past_kv from previous generation\n",
+        "\n",
+        "            print(f\"Full text so far: {full_text}\")\n",
+        "            print(\"-\" * 50)\n",
+        "\n",
+        "        print(\"\\n==== Final Story ====\")\n",
+        "        print(full_text)\n",
+        "\n",
+        "        return full_text\n",
+        "\n",
+        "    # 4. Run the test\n",
+        "    try:\n",
+        "        print(\"\\n=== Testing Manual KV Cache Usage ===\\n\")\n",
+        "\n",
+        "        story = run_manual_multi_turn_generation()\n",
+        "\n",
+        "        print(\"\\n=== Test Complete ===\")\n",
+        "\n",
+        "    except Exception as e:\n",
+        "        import traceback\n",
+        "        print(f\"Error during testing: {e}\")\n",
+        "        traceback.print_exc()\n",
+        "\n",
+        "# Run the test\n",
+        "test_direct_kv_cache_usage()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "p2_BiDS0IlPl",
+        "outputId": "f62d0e49-8f27-4c9b-aeac-dbe1fe1def21"
+      },
+      "execution_count": 3,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Using device: cuda\n",
+            "\n",
+            "=== Testing Manual KV Cache Usage ===\n",
+            "\n",
+            "\n",
+            "==== Turn 1 ====\n",
+            "Current input: Once upon a time, in a magical forest, there lived a\n",
+            "KV cache present: False\n",
+            "Press Enter to generate next part...\n",
+            "Generated:  wizard and a witch. One day, they were attacked by a flying beast and the wizard fled to\n",
+            "\n",
+            "==== Injecting custom text:  Suddenly, a rainbow appeared in the sky! ====\n",
+            "Full text so far: Once upon a time, in a magical forest, there lived a wizard and a witch. One day, they were attacked by a flying beast and the wizard fled to Suddenly, a rainbow appeared in the sky!\n",
+            "--------------------------------------------------\n",
+            "\n",
+            "==== Turn 2 ====\n",
+            "Current input:  Suddenly, a rainbow appeared in the sky!\n",
+            "KV cache present: True\n",
+            "Press Enter to generate next part...\n",
+            "Generated:  Although it was a little less than a hundred meters wide, it was enough to cover the entire sky\n",
+            "\n",
+            "==== Injecting custom text:  Suddenly, a rainbow appeared in the sky! ====\n",
+            "Full text so far: Once upon a time, in a magical forest, there lived a wizard and a witch. One day, they were attacked by a flying beast and the wizard fled to Suddenly, a rainbow appeared in the sky! Although it was a little less than a hundred meters wide, it was enough to cover the entire sky Suddenly, a rainbow appeared in the sky!\n",
+            "--------------------------------------------------\n",
+            "\n",
+            "==== Turn 3 ====\n",
+            "Current input:  Suddenly, a rainbow appeared in the sky!\n",
+            "KV cache present: True\n",
+            "Press Enter to generate next part...\n",
+            "Generated:  A rainbow!\n",
+            "\n",
+            "I thought about it, and after seeing how many people were in the sky\n",
+            "\n",
+            "==== Injecting custom text:  Suddenly, a rainbow appeared in the sky! ====\n",
+            "Full text so far: Once upon a time, in a magical forest, there lived a wizard and a witch. One day, they were attacked by a flying beast and the wizard fled to Suddenly, a rainbow appeared in the sky! Although it was a little less than a hundred meters wide, it was enough to cover the entire sky Suddenly, a rainbow appeared in the sky! A rainbow!\n",
+            "\n",
+            "I thought about it, and after seeing how many people were in the sky Suddenly, a rainbow appeared in the sky!\n",
+            "--------------------------------------------------\n",
+            "\n",
+            "==== Final Story ====\n",
+            "Once upon a time, in a magical forest, there lived a wizard and a witch. One day, they were attacked by a flying beast and the wizard fled to Suddenly, a rainbow appeared in the sky! Although it was a little less than a hundred meters wide, it was enough to cover the entire sky Suddenly, a rainbow appeared in the sky! A rainbow!\n",
+            "\n",
+            "I thought about it, and after seeing how many people were in the sky Suddenly, a rainbow appeared in the sky!\n",
+            "\n",
+            "=== Test Complete ===\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def test_retool_with_working_kv_cache():\n",
+        "    # 1. Setup model, tokenizer, and device\n",
+        "    from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+        "    import torch\n",
+        "    import re\n",
+        "\n",
+        "    # Use a model that fits in memory\n",
+        "    model_name = \"gpt2-medium\"\n",
+        "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
+        "\n",
+        "    # Check device\n",
+        "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "    print(f\"Using device: {device}\")\n",
+        "\n",
+        "    # Load model to device\n",
+        "    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)\n",
+        "\n",
+        "    # 2. Add special tokens\n",
+        "    special_tokens = {\n",
+        "        'additional_special_tokens': ['<code>', '</code>', '<interpreter>', '</interpreter>']\n",
+        "    }\n",
+        "    tokenizer.add_special_tokens(special_tokens)\n",
+        "    model.resize_token_embeddings(len(tokenizer))\n",
+        "\n",
+        "    # Get token IDs\n",
+        "    code_start_id = tokenizer.convert_tokens_to_ids('<code>')\n",
+        "    code_end_id = tokenizer.convert_tokens_to_ids('</code>')\n",
+        "    interpreter_start_id = tokenizer.convert_tokens_to_ids('<interpreter>')\n",
+        "    interpreter_end_id = tokenizer.convert_tokens_to_ids('</interpreter>')\n",
+        "\n",
+        "    print(f\"EOS token ID: {tokenizer.eos_token_id}\")\n",
+        "    print(f\"Code tokens: {code_start_id}, {code_end_id}\")\n",
+        "    print(f\"Interpreter tokens: {interpreter_start_id}, {interpreter_end_id}\")\n",
+        "\n",
+        "    # 3. Manual token generation with KV caching\n",
+        "    def generate_with_manual_kv_cache(input_ids, past_key_values=None, max_tokens=20, stop_ids=None):\n",
+        "        \"\"\"Generate tokens with KV cache until a stop token or max_tokens is reached\"\"\"\n",
+        "        if stop_ids is None:\n",
+        "            stop_ids = [tokenizer.eos_token_id]\n",
+        "\n",
+        "        current_ids = input_ids.clone()\n",
+        "        generated_tokens = []\n",
+        "\n",
+        "        for _ in range(max_tokens):\n",
+        "            # Forward pass with past_key_values\n",
+        "            with torch.no_grad():\n",
+        "                outputs = model(\n",
+        "                    input_ids=current_ids if past_key_values is None else current_ids[:, -1:],\n",
+        "                    past_key_values=past_key_values,\n",
+        "                    use_cache=True\n",
+        "                )\n",
+        "\n",
+        "            # Get logits for the next token\n",
+        "            next_token_logits = outputs.logits[:, -1, :]\n",
+        "\n",
+        "            # Sample from the distribution\n",
+        "            probs = torch.nn.functional.softmax(next_token_logits / 0.7, dim=-1)\n",
+        "            next_token = torch.multinomial(probs, num_samples=1)\n",
+        "\n",
+        "            # Get the token ID\n",
+        "            token_id = next_token.item()\n",
+        "\n",
+        "            # Add to generated tokens\n",
+        "            generated_tokens.append(token_id)\n",
+        "\n",
+        "            # Update current_ids for next iteration\n",
+        "            current_ids = torch.cat([current_ids, next_token], dim=1)\n",
+        "\n",
+        "            # Update past_key_values\n",
+        "            past_key_values = outputs.past_key_values\n",
+        "\n",
+        "            # Check if we hit a stop token\n",
+        "            if token_id in stop_ids:\n",
+        "                break\n",
+        "\n",
+        "        # Convert list of token IDs to tensor\n",
+        "        result_tensor = torch.tensor([generated_tokens], device=device)\n",
+        "        return result_tensor, past_key_values\n",
+        "\n",
+        "    # 4. ReTool simulation with working KV cache\n",
+        "    def simulate_retool_with_working_kv_cache(prompt, max_turns=3):\n",
+        "        \"\"\"Simulate the ReTool process with working KV cache\"\"\"\n",
+        "        # Tokenize the prompt\n",
+        "        prompt_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
+        "\n",
+        "        # Initialize tracking\n",
+        "        full_sequence = prompt_ids.clone()\n",
+        "        completion = torch.empty((1, 0), dtype=torch.long, device=device)\n",
+        "        interpreter_positions = []\n",
+        "\n",
+        "        # Keep the KV cache from previous turns\n",
+        "        past_kv = None\n",
+        "\n",
+        "        for turn_idx in range(max_turns):\n",
+        "            print(f\"\\n==== Turn {turn_idx + 1} ====\")\n",
+        "\n",
+        "            # Determine what to generate from\n",
+        "            if turn_idx == 0:\n",
+        "                # First turn - generate from the prompt\n",
+        "                current_input = full_sequence\n",
+        "                print(f\"Generating from prompt: {tokenizer.decode(current_input[0])}\")\n",
+        "            else:\n",
+        "                # Later turns - might be generating from interpreter output\n",
+        "                current_input = full_sequence[:, -20:] if full_sequence.size(1) > 20 else full_sequence\n",
+        "                print(f\"Generating from: {tokenizer.decode(current_input[0])}\")\n",
+        "\n",
+        "            # Generate with manual KV cache\n",
+        "            new_tokens, past_kv = generate_with_manual_kv_cache(\n",
+        "                current_input,\n",
+        "                past_key_values=past_kv,\n",
+        "                max_tokens=30,\n",
+        "                stop_ids=[tokenizer.eos_token_id, code_end_id]\n",
+        "            )\n",
+        "\n",
+        "            # Decode and display\n",
+        "            new_text = tokenizer.decode(new_tokens[0])\n",
+        "            print(f\"Generated: {new_text}\")\n",
+        "\n",
+        "            # Update tracking\n",
+        "            full_sequence = torch.cat([full_sequence, new_tokens], dim=1)\n",
+        "            completion = torch.cat([completion, new_tokens], dim=1)\n",
+        "\n",
+        "            # Check for code blocks\n",
+        "            full_text = tokenizer.decode(full_sequence[0])\n",
+        "            code_blocks = re.findall(r'<code>(.*?)</code>', full_text, re.DOTALL)\n",
+        "\n",
+        "            # Pause for inspection\n",
+        "            input(\"Press Enter to continue...\")\n",
+        "\n",
+        "            if code_blocks and code_end_id in new_tokens[0]:\n",
+        "                print(\"\\n==== Found code block! ====\")\n",
+        "                # Get the last code block\n",
+        "                code_block = code_blocks[-1].strip()\n",
+        "                print(f\"Code block: {code_block}\")\n",
+        "\n",
+        "                # Mock code execution\n",
+        "                print(\"\\n==== Executing code ====\")\n",
+        "                interpreter_output = \"0 1 1 2 3\"\n",
+        "                print(f\"Execution result: {interpreter_output}\")\n",
+        "\n",
+        "                # Format interpreter feedback\n",
+        "                interpreter_text = f\"<interpreter>{interpreter_output}</interpreter>\"\n",
+        "                interpreter_ids = tokenizer.encode(\n",
+        "                    interpreter_text,\n",
+        "                    return_tensors=\"pt\",\n",
+        "                    add_special_tokens=False\n",
+        "                ).to(device)\n",
+        "\n",
+        "                # Record positions\n",
+        "                start_idx = completion.size(1)\n",
+        "                completion = torch.cat([completion, interpreter_ids], dim=1)\n",
+        "                end_idx = completion.size(1) - 1\n",
+        "                interpreter_positions.append((start_idx, end_idx))\n",
+        "\n",
+        "                # Add to full sequence\n",
+        "                full_sequence = torch.cat([full_sequence, interpreter_ids], dim=1)\n",
+        "                print(f\"Added interpreter output: {interpreter_text}\")\n",
+        "\n",
+        "                # We're still using the same past_kv for the next turn\n",
+        "                # The next input will be the interpreter output\n",
+        "            elif tokenizer.eos_token_id in new_tokens[0]:\n",
+        "                print(\"Found EOS token, ending generation\")\n",
+        "                break\n",
+        "\n",
+        "        return completion, interpreter_positions\n",
+        "\n",
+        "    # 5. Test with a prompt containing a code block\n",
+        "    prompt = \"\"\"Let me solve this problem with code:\n",
+        "\n",
+        "<code>\n",
+        "def fibonacci(n):\n",
+        "    a, b = 0, 1\n",
+        "    result = []\n",
+        "    for _ in range(n):\n",
+        "        result.append(a)\n",
+        "        a, b = b, a + b\n",
+        "    return result\n",
+        "\n",
+        "print(fibonacci(5))\n",
+        "</code>\"\"\"\n",
+        "\n",
+        "    # 6. Run the test\n",
+        "    try:\n",
+        "        print(\"\\n=== Testing ReTool with Working KV Cache ===\\n\")\n",
+        "\n",
+        "        completion, positions = simulate_retool_with_working_kv_cache(prompt)\n",
+        "\n",
+        "        print(\"\\n=== Final Results ===\\n\")\n",
+        "        print(\"Generated completion:\")\n",
+        "        print(tokenizer.decode(completion[0]))\n",
+        "\n",
+        "        print(\"\\nFull text:\")\n",
+        "        print(tokenizer.decode(torch.cat([tokenizer.encode(prompt, return_tensors=\"pt\")[0].to(device), completion[0]])))\n",
+        "\n",
+        "        print(\"\\nInterpreter positions:\", positions)\n",
+        "\n",
+        "    except Exception as e:\n",
+        "        import traceback\n",
+        "        print(f\"Error during testing: {e}\")\n",
+        "        traceback.print_exc()\n",
+        "\n",
+        "# Run the test\n",
+        "test_retool_with_working_kv_cache()"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "T6_ob3S4M5mn",
+        "outputId": "e5f42a03-c49a-403f-d27b-0ae50ecd095e"
+      },
+      "execution_count": 4,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Using device: cuda\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "The new embeddings will be initialized from a multivariate normal distribution that has old embeddings' mean and covariance. As described in this article: https://nlp.stanford.edu/~johnhew/vocab-expansion.html. To disable this, use `mean_resizing=False`\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "EOS token ID: 50256\n",
+            "Code tokens: 50257, 50258\n",
+            "Interpreter tokens: 50259, 50260\n",
+            "\n",
+            "=== Testing ReTool with Working KV Cache ===\n",
+            "\n",
+            "\n",
+            "==== Turn 1 ====\n",
+            "Generating from prompt: Let me solve this problem with code:\n",
+            "\n",
+            "<code>\n",
+            "def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "</code>\n",
+            "Generated: \n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = [0,\n",
+            "Press Enter to continue...\n",
+            "\n",
+            "==== Turn 2 ====\n",
+            "Generating from:     a, b = 0, 1\n",
+            "\n",
+            "    result = [0,\n",
+            "Generated:  0, 0, 1]\n",
+            "\n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "\n",
+            "ret = [0,\n",
+            "Press Enter to continue...\n",
+            "\n",
+            "==== Turn 3 ====\n",
+            "Generating from:      a, b = b, a + b\n",
+            "\n",
+            "\n",
+            "ret = [0,\n",
+            "Generated:  1, 1, 1]\n",
+            "\n",
+            "for i,j in enumerate(n, fibonacci(n-1, 1-f\n",
+            "Press Enter to continue...\n",
+            "\n",
+            "=== Final Results ===\n",
+            "\n",
+            "Generated completion:\n",
+            "\n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = [0, 0, 0, 1]\n",
+            "\n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "\n",
+            "ret = [0, 1, 1, 1]\n",
+            "\n",
+            "for i,j in enumerate(n, fibonacci(n-1, 1-f\n",
+            "\n",
+            "Full text:\n",
+            "Let me solve this problem with code:\n",
+            "\n",
+            "<code>\n",
+            "def fibonacci(n):\n",
+            "    a, b = 0, 1\n",
+            "    result = []\n",
+            "    for _ in range(n):\n",
+            "        result.append(a)\n",
+            "        a, b = b, a + b\n",
+            "    return result\n",
+            "\n",
+            "print(fibonacci(5))\n",
+            "</code>\n",
+            "def fibonacci(n):\n",
+            "\n",
+            "    a, b = 0, 1\n",
+            "\n",
+            "    result = [0, 0, 0, 1]\n",
+            "\n",
+            "       a, b = b, a + b\n",
+            "\n",
+            "\n",
+            "ret = [0, 1, 1, 1]\n",
+            "\n",
+            "for i,j in enumerate(n, fibonacci(n-1, 1-f\n",
+            "\n",
+            "Interpreter positions: []\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "YFIXEa5fM5px"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "FjaszXJOIlVz"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "xgjX6_xZaCDQ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "iTGXE8lRaCF4"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "oM5BSZHEaCIx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7d252539"
+      },
+      "source": [
+        "**1. Clear CUDA Cache:**\n",
+        "\n",
+        "This is often the first thing to try when you get a CUDA OOM error."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "YhKSjnxiaBCb"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "f793cb16",
+        "outputId": "3b5b2b99-2e9b-44a2-88df-7293e51de014"
+      },
+      "source": [
+        "import torch\n",
+        "\n",
+        "if torch.cuda.is_available():\n",
+        "    torch.cuda.empty_cache()\n",
+        "    print(\"CUDA cache cleared!\")\n",
+        "else:\n",
+        "    print(\"CUDA not available, no cache to clear.\")"
+      ],
+      "execution_count": 18,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "CUDA cache cleared!\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d25e30fe"
+      },
+      "source": [
+        "**2. Delete Large Variables and Run Garbage Collection:**\n",
+        "\n",
+        "Identify variables holding large objects (like models, tensors, dataframes) that you don't need anymore and delete them. Then explicitly run garbage collection."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "02474dce",
+        "outputId": "80223089-31f7-485f-8490-aad00d97277a"
+      },
+      "source": [
+        "# Example: if you have a large model or tensor named 'model' or 'data'\n",
+        "# del model\n",
+        "# del data\n",
+        "\n",
+        "import gc\n",
+        "gc.collect()\n",
+        "\n",
+        "print(\"Garbage collection complete.\")"
+      ],
+      "execution_count": 19,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Garbage collection complete.\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "105cefce"
+      },
+      "source": [
+        "**3. Restart Runtime:**\n",
+        "\n",
+        "If the above steps don't work, restarting the runtime is the most drastic but often most effective way to clear all memory. Go to the Colab menu: `Runtime` -> `Restart runtime`."
+      ]
+    }
+  ]
+}