apply quentin's update here

bonus-unit1/bonus-unit1.ipynb

@@ -184,7 +184,7 @@
     "import torch\n",
     "import json\n",
     "\n",
-    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig, set_seed\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed\n",
     "from datasets import load_dataset\n",
     "from trl import SFTConfig, SFTTrainer\n",
     "from peft import LoraConfig, TaskType\n",
@@ -319,10 +319,7 @@
    "source": [
     "dataset = dataset.map(preprocess, remove_columns=\"messages\")\n",
     "dataset = dataset[\"train\"].train_test_split(0.1)\n",
-    "print(dataset)\n",
-    "\n",
-    "dataset[\"train\"] = dataset[\"train\"].select(range(100))\n",
-    "dataset[\"test\"] = dataset[\"test\"].select(range(10))"
+    "print(dataset)"
    ]
   },
   {