diff --git a/mlx/lora/finetune-experiments.ipynb b/mlx/lora/finetune-experiments.ipynb
index 1b6242e50b0a262e35b5ab2e4a248721d51ec0ab..d71c2963bb6682ddc25e7f773a147a669ecd4eb9 100644
--- a/mlx/lora/finetune-experiments.ipynb
+++ b/mlx/lora/finetune-experiments.ipynb
@@ -79,19 +79,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 35,
    "outputs": [],
    "source": [
     "system_message =\"\"\"\n",
     "You are a text processing agent. As instructed below, extract information from the provided content in a structured format without discussing reasoning or providing commentary. Only use source text given as input for data extraction unless specifically asked for inference.\n",
-    "\"\"\"\n",
+    "\"\"\".strip()\n",
     "\n",
     "instruction = \"\"\"\n",
-    "Analyze the content from a German law journal's website which follows this instruction. Your task is to identify members of the editorial board (terms to look for: 'Herausgeber', 'Redakteur', 'Schriftleitung') and the advisory board ('Beirat'). For each identified member, extract and organize their information into the following categories: lastname, firstname, title (including academic titles like 'Dr.' or 'Prof. Dr.' and suffixes such as 'LL.M.'), position (their job title, if provided), affiliation, and role. For 'role', infer the role within the journal from the context (options 'Herausgeber', 'Redaktion', 'Schriftleitung', 'Beirat', or an empty string if the role is unknown).\n",
+    "Analyze content from a German law journal's website. Your task is to identify members of the editorial board (terms to look for: 'Herausgeber', 'Redakteur', 'Schriftleitung') and the advisory board ('Beirat'). For each identified member, extract and organize their information into the following categories: lastname, firstname, title (including academic titles like 'Dr.' or 'Prof. Dr.' and suffixes such as 'LL.M.'), position (their job title, if provided), affiliation, and role. For 'role', infer the role within the journal from the context (options 'Herausgeber', 'Redaktion', 'Schriftleitung', 'Beirat', or an empty string if the role is unknown).\n",
     "\n",
     "- Format the output as a YAML list of dictionaries.\n",
     "- Exclude any dictionary entries for which information is not available or relevant fields are empty.\n",
-    "- Ensure the YAML output is strictly valid.\n",
+    "- Ensure the YAML output is strictly valid. It must be a list of dictionaries. \n",
+    "\n",
+    "Here is an example:\n",
+    "\n",
+    "```yaml\n",
+    "- lastname: Mustermann\n",
+    "  firstname: Martina\n",
+    "  title: Dr.\n",
+    "  position: Vorsitzender Richterin\n",
+    "  affiliation: Oberlandesgericht Buxtehude\n",
+    "  role: Herausgeber\n",
+    "```\n",
     "\n",
     "Adhere to these guidelines to efficiently and accurately process the following content:\"\n",
     "\"\"\".strip()"
@@ -99,8 +110,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-29T14:10:39.997135Z",
-     "start_time": "2024-02-29T14:10:39.990995Z"
+     "end_time": "2024-02-29T17:50:20.348433Z",
+     "start_time": "2024-02-29T17:50:20.344221Z"
     }
    },
    "id": "b4be7c0872d2fd34"
@@ -364,16 +375,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 36,
    "outputs": [],
    "source": [
     "prompt=f\"\"\"\n",
+    "### SYSTEM\n",
     "{system_message}\n",
-    "\n",
+    "### USER\n",
     "{instruction}\n",
-    "\n",
     "### CONTENT\n",
-    "\n",
     "Herausgeber:\n",
     "Prof. Dr. Stefan Knesebeck, Universität Wuppertal\n",
     "Prof. Dr. Dr. h.c. Fritz M. Müller LL.M.(Yale), Universität Wanne-Eickel\n",
@@ -383,40 +393,48 @@
     "Redaktion:\n",
     "RA Adam Gengelbach, Unterhachingen\n",
     "Ass. iur. Petra Priem, Herrenchiemsee\n",
-    "\n",
     "### END OF CONTENT\n",
     "\"\"\".strip()"
    ],
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-29T14:16:06.332197Z",
-     "start_time": "2024-02-29T14:16:06.326960Z"
+     "end_time": "2024-02-29T17:50:29.480885Z",
+     "start_time": "2024-02-29T17:50:29.476367Z"
     }
    },
    "id": "8d316a1e7570f1d4"
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 37,
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "### SYSTEM\n",
       "You are a text processing agent. As instructed below, extract information from the provided content in a structured format without discussing reasoning or providing commentary. Only use source text given as input for data extraction unless specifically asked for inference.\n",
-      "\n",
-      "\n",
+      "### USER\n",
       "Analyze the content from a German law journal's website which follows this instruction. Your task is to identify members of the editorial board (terms to look for: 'Herausgeber', 'Redakteur', 'Schriftleitung') and the advisory board ('Beirat'). For each identified member, extract and organize their information into the following categories: lastname, firstname, title (including academic titles like 'Dr.' or 'Prof. Dr.' and suffixes such as 'LL.M.'), position (their job title, if provided), affiliation, and role. For 'role', infer the role within the journal from the context (options 'Herausgeber', 'Redaktion', 'Schriftleitung', 'Beirat', or an empty string if the role is unknown).\n",
       "\n",
       "- Format the output as a YAML list of dictionaries.\n",
       "- Exclude any dictionary entries for which information is not available or relevant fields are empty.\n",
-      "- Ensure the YAML output is strictly valid.\n",
+      "- Ensure the YAML output is strictly valid. It must be a list of dictionaries. \n",
       "\n",
-      "Adhere to these guidelines to efficiently and accurately process the following content:\"\n",
+      "Here is an example:\n",
       "\n",
-      "### CONTENT\n",
+      "```yaml\n",
+      "- lastname: Mustermann\n",
+      "  firstname: Martina\n",
+      "  title: Dr.\n",
+      "  position: Vorsitzender Richterin\n",
+      "  affiliation: Oberlandesgericht Buxtehude\n",
+      "  role: Herausgeber\n",
+      "```\n",
       "\n",
+      "Adhere to these guidelines to efficiently and accurately process the following content:\"\n",
+      "### CONTENT\n",
       "Herausgeber:\n",
       "Prof. Dr. Stefan Knesebeck, Universität Wuppertal\n",
       "Prof. Dr. Dr. h.c. Fritz M. Müller LL.M.(Yale), Universität Wanne-Eickel\n",
@@ -426,7 +444,6 @@
       "Redaktion:\n",
       "RA Adam Gengelbach, Unterhachingen\n",
       "Ass. iur. Petra Priem, Herrenchiemsee\n",
-      "\n",
       "### END OF CONTENT\n"
      ]
     }
@@ -437,15 +454,15 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-29T14:16:07.677659Z",
-     "start_time": "2024-02-29T14:16:07.672036Z"
+     "end_time": "2024-02-29T17:50:30.614801Z",
+     "start_time": "2024-02-29T17:50:30.609304Z"
     }
    },
    "id": "3e7a823a9f4a35d9"
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 38,
    "outputs": [
     {
      "name": "stdout",
@@ -453,40 +470,48 @@
      "text": [
       "\r\n",
       "\r\n",
-      "### OUTPUT\r\n",
-      "\r\n",
+      "``\r\n",
       "- lastname: Knesebeck\r\n",
       "  firstname: Stefan\r\n",
       "  title: Prof. Dr.\r\n",
-      "  position: Herausgeber\r\n",
+      "  position: Vorsitzender Richter\r\n",
       "  affiliation: Universität Wuppertal\r\n",
       "  role: Herausgeber\r\n",
       "- lastname: Müller\r\n",
       "  firstname: Fritz M.\r\n",
-      "  title: Prof. Dr. Dr. h.c.\r\n",
-      "  position: Herausgeber\r\n",
+      "  title: Prof. Dr. Dr. h.c. LL.M.(Yale)\r\n",
+      "  position: Vorsitzender Richter\r\n",
       "  affiliation: Universität Wanne-Eickel\r\n",
       "  role: Herausgeber\r\n",
       "- lastname: Valentin\r\n",
       "  firstname: Vera\r\n",
-      "  title: RA Prof. Dr.\r\n",
-      "  position: Redaktion\r\n",
+      "  title: Prof. Dr.\r\n",
+      "  position: Vorsitzenderin\r\n",
       "  affiliation: Hochschule für Recht und Sport Edingen\r\n",
-      "  role: Redaktion\r\n",
+      "  role: Herausgeber\r\n",
       "- lastname: Rosenbaum\r\n",
       "  firstname: Rita\r\n",
       "  title: Prof. Dr. Dr. h.c.\r\n",
-      "  position: Herausgeber\r\n",
+      "  position: Vorsitzenderin\r\n",
       "  affiliation: Universität Tupfingen\r\n",
       "  role: Herausgeber\r\n",
       "- lastname: Gonzalo de Sanchez\r\n",
       "  firstname: Ingo\r\n",
       "  title: Dr.\r\n",
-      "  position: Vorsitzender Richter am Oberlandesgericht Rostock\r\n",
-      "  affiliation: Empty\r\n",
-      "  role: Empty\r\n",
-      "\r\n",
-      "Generation took 65.16657900810242 seconds\n"
+      "  position: Vorsitzender Richter am Oberlandesgericht\r\n",
+      "  affiliation: Rostock\r\n",
+      "  role: Herausgeber\r\n",
+      "- lastname: Gengelbach\r\n",
+      "  firstname: Adam\r\n",
+      "  position: RA\r\n",
+      "  affiliation: Unterhachingen\r\n",
+      "  role: Redaktion\r\n",
+      "- lastname: Priem\r\n",
+      "  firstname: Petra\r\n",
+      "  position: Ass. iur.\r\n",
+      "  affiliation: Herrenchiemsee\r\n",
+      "  role: Redaktion\r\n",
+      "Generation took 79.5884530544281 seconds\n"
      ]
     }
    ],
@@ -507,8 +532,8 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-29T14:17:15.637298Z",
-     "start_time": "2024-02-29T14:16:10.463466Z"
+     "end_time": "2024-02-29T17:51:54.067151Z",
+     "start_time": "2024-02-29T17:50:34.468241Z"
     }
    },
    "id": "1ea4b39f35c09268"
@@ -516,7 +541,7 @@
   {
    "cell_type": "markdown",
    "source": [
-    "# mlx-community/quantized-gemma-7b-it"
+    "## mlx-community/quantized-gemma-7b-it"
    ],
    "metadata": {
     "collapsed": false
@@ -525,7 +550,52 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 46,
+   "outputs": [],
+   "source": [
+    "system_message =\"\"\"\n",
+    "You are a text processing agent. As instructed below, extract information from the provided content in a structured format without discussing reasoning or providing commentary. Only use source text given as input for data extraction unless specifically asked for inference.\n",
+    "\"\"\".strip()\n",
+    "\n",
+    "instruction = \"\"\"\n",
+    "Analyze content from a German law journal's website. Your task is to identify members of the editorial board (terms to look for: 'Herausgeber', 'Redakteur', 'Schriftleitung') and the advisory board ('Beirat'). For each identified member, extract and organize their information into the following categories: lastname, firstname, title (including academic titles like 'Dr.' or 'Prof. Dr.' and suffixes such as 'LL.M.'), position (their job title, if provided), affiliation, and role. For 'role', infer the role within the journal from the context (options 'Herausgeber', 'Redaktion', 'Schriftleitung', 'Beirat', or an empty string if the role is unknown).\n",
+    "Format of the output:\n",
+    "- Format the output as a YAML list of dictionaries.\n",
+    "- Exclude any dictionary entries for which information is not available or relevant fields are empty.\n",
+    "- Ensure the YAML output is strictly valid. It must be a list of dictionaries. Return only the YAML and nothing else.\n",
+    "\"\n",
+    "\"\"\".strip()\n",
+    "\n",
+    "prompt=f\"\"\"\n",
+    "### INSTRUCTION\n",
+    "{system_message}\n",
+    "### USER\n",
+    "{instruction}\n",
+    "### CONTENT\n",
+    "Herausgeber:\n",
+    "Prof. Dr. Stefan Knesebeck, Universität Wuppertal\n",
+    "Prof. Dr. Dr. h.c. Fritz M. Müller LL.M.(Yale), Universität Wanne-Eickel\n",
+    "RA Prof. Dr. Vera Valentin, Hochschule für Recht und Sport Edingen\n",
+    "Prof. Dr. Dr. h.c. Rita Rosenbaum, Universität Tupfingen\n",
+    "Dr. Ingo Gonzalo de Sanchez, Vorsitzender Richter am Oberlandesgericht Rostock\n",
+    "Redaktion:\n",
+    "RA Adam Gengelbach, Unterhachingen\n",
+    "Ass. iur. Petra Priem, Herrenchiemsee\n",
+    "### END OF CONTENT\n",
+    "\"\"\".strip()"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-02-29T20:07:53.541433Z",
+     "start_time": "2024-02-29T20:07:53.535558Z"
+    }
+   },
+   "id": "485a8723b4668dc7"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
    "outputs": [
     {
      "data": {
@@ -533,7 +603,7 @@
       "application/vnd.jupyter.widget-view+json": {
        "version_major": 2,
        "version_minor": 0,
-       "model_id": "a229bb7e24b3445fa000e2b8f9b6ec81"
+       "model_id": "7dbafccfbe7242f4aed67cbc17ee4a49"
       }
      },
      "metadata": {},
@@ -543,58 +613,58 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "\n",
+      "\n",
+      "**Expected Output:**\n",
+      "\n",
       "```yaml\n",
       "- lastname: Knesebeck\n",
       "  firstname: Stefan\n",
       "  title: Prof. Dr.\n",
       "  position:\n",
       "  affiliation: Universität Wuppertal\n",
-      "  role:Herausgeber\n",
+      "  role: Herausgeber\n",
       "\n",
       "- lastname: Müller\n",
       "  firstname: Dr. Dr. h.c. Fritz M.\n",
       "  title: Prof. Dr. LL.M.(Yale)\n",
       "  position:\n",
       "  affiliation: Universität Wanne-Eickel\n",
-      "  role:Herausgeber\n",
+      "  role: Herausgeber\n",
       "\n",
       "- lastname: Valentin\n",
-      "  firstname: RA Prof. Dr. Vera\n",
-      "  title: Prof. Dr.\n",
+      "  firstname: Prof. Dr. Vera\n",
       "  position:\n",
       "  affiliation: Hochschule für Recht und Sport Edingen\n",
-      "  role:Herausgeber\n",
+      "  role: Herausgeber\n",
       "\n",
       "- lastname: Rosenbaum\n",
       "  firstname: Prof. Dr. Dr. h.c. Rita\n",
-      "  title: Prof. Dr.\n",
       "  position:\n",
       "  affiliation: Universität Tupfingen\n",
-      "  role:Herausgeber\n",
+      "  role: Herausgeber\n",
       "\n",
       "- lastname: Gonzalo de Sanchez\n",
       "  firstname: Dr. Ingo\n",
-      "  title:\n",
-      "  position: Vorsitzender Richter am Oberlandesgericht Rostock\n",
+      "  position:\n",
+      "  affiliation: Oberlandesgericht Rostock\n",
       "  role:\n",
       "\n",
       "- lastname: Gengelbach\n",
       "  firstname: RA Adam\n",
-      "  title:\n",
       "  position:\n",
       "  affiliation: Unterhachingen\n",
-      "  role:\n",
+      "  role: Redakteur\n",
       "\n",
       "- lastname: Priem\n",
       "  firstname: Ass. iur. Petra\n",
-      "  title:\n",
       "  position:\n",
       "  affiliation: Herrenchiemsee\n",
-      "  role:\n",
+      "  role: Redakteur\n",
       "```\n",
       "\n",
-      "This is the requested output. Please extract the requested\n",
-      "Generation took 35.79903483390808 seconds\n"
+      "**Note:** The text does not specify the role of Dr. Gonzalo de Sanchez, therefore the 'role' field for this member is empty\n",
+      "Generation took 51.236616134643555 seconds\n"
      ]
     }
    ],
@@ -605,15 +675,15 @@
     "\n",
     "model, tokenizer = load(\"mlx-community/quantized-gemma-7b-it\")\n",
     "start_time = time.time()\n",
-    "response = generate(model, tokenizer, prompt=prompt, verbose=False, max_tokens=300)\n",
+    "response = generate(model, tokenizer, prompt=prompt, verbose=False, max_tokens=300, temp=0)\n",
     "print(response)\n",
     "print(f'Generation took {time.time() - start_time} seconds')\n"
    ],
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-29T08:38:10.003959Z",
-     "start_time": "2024-02-29T08:37:23.119783Z"
+     "end_time": "2024-02-29T20:08:54.748222Z",
+     "start_time": "2024-02-29T20:07:54.439115Z"
     }
    },
    "id": "89e1a05fc3b6e435"
@@ -622,7 +692,9 @@
    "cell_type": "code",
    "execution_count": null,
    "outputs": [],
-   "source": [],
+   "source": [
+    "{\"text\": \"<bos><start_of_turn>user\\nWhat is the capital of France?<end_of_turn>\\n<start_of_turn>model\\nParis is the capital of France.<end_of_turn><eos>\"}"
+   ],
    "metadata": {
     "collapsed": false
    },