embedding_finetune.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "from transformers import AutoTokenizer, AutoModel, AutoConfig\n",
    "import torch.optim as optim\n",
    "from mlx_lm import load, generate\n",
    "from tqdm import tqdm\n",
    "import chromadb\n",
    "from chromadb.utils import embedding_functions\n",
    "import torch\n",
    "from chromadb import Documents, EmbeddingFunction, Embeddings\n",
    "from numpy.typing import ArrayLike\n",
    "from sentence_transformers.util import cos_sim, dot_score\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "from peft import PeftModel\n",
    "from llm2vec import LLM2Vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Embedding(EmbeddingFunction):\n",
    "    def __init__(self):\n",
    "        self.embedding_model = SentenceTransformer('thenlper/gte-large')  # (\"Salesforce/SFR-Embedding-Mistral\") (7b)\n",
    "\n",
    "    def encode(self, input: Documents) -> Embeddings:\n",
    "        if isinstance(input, str):\n",
    "            input = [input]\n",
    "        embeddings = self.embedding_model.encode(input)\n",
    "        return embeddings\n",
    "    \n",
    "    def __call__(self, input: Documents) -> Embeddings:\n",
    "        if isinstance(input, str):\n",
    "            input = [input]\n",
    "        embeddings = self.embedding_model.encode(input)\n",
    "        return embeddings\n",
    "\n",
    "    def distance(self, x: ArrayLike, y: ArrayLike) -> float:\n",
    "        return cos_sim(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PeftEmbeddingModel(EmbeddingFunction):\n",
    "    def __init__(self):\n",
    "        config = AutoConfig.from_pretrained(\"McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp\", trust_remote_code=True)\n",
    "        tokenizer = AutoTokenizer.from_pretrained(\"McGill-NLP/LLM2Vec-Sheared-LLaMA-mntp\")\n",
    "\n",
    "        # Loading base MNTP model, along with custom code that enables bidirectional connections in decoder-only LLMs\n",
    "        tokenizer = AutoTokenizer.from_pretrained(\n",
    "            \"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\"\n",
    "        )\n",
    "        config = AutoConfig.from_pretrained(\n",
    "            \"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\", trust_remote_code=True\n",
    "        )\n",
    "        model = AutoModel.from_pretrained(\n",
    "            \"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\",\n",
    "            trust_remote_code=True,\n",
    "            config=config,\n",
    "            torch_dtype=torch.float16,\n",
    "            device_map=\"mps\",\n",
    "        )\n",
    "        model = PeftModel.from_pretrained(\n",
    "            model,\n",
    "            \"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\",\n",
    "        )\n",
    "        model = model.merge_and_unload()  # This can take several minutes on cpu\n",
    "\n",
    "        # Loading unsupervised-trained LoRA weights. This loads the trained LoRA weights on top of MNTP model. Hence the final weights are -- Base model + MNTP (LoRA) + SimCSE (LoRA).\n",
    "        model = PeftModel.from_pretrained(\n",
    "            model, \"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-unsup-simcse\"\n",
    "        )\n",
    "\n",
    "        # Or loading supervised-trained LoRA weights\n",
    "        model = PeftModel.from_pretrained(\n",
    "            model, \"McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-supervised\"\n",
    "        )\n",
    "        self.l2v = LLM2Vec(model, tokenizer, pooling_mode=\"mean\", max_length=512)\n",
    "\n",
    "    def __call__(self, input: Documents) -> Embeddings:\n",
    "        if isinstance(input, str):\n",
    "            input = [input]\n",
    "        embeddings = self.l2v.encode(input)\n",
    "        return embeddings\n",
    "    \n",
    "    def encode(self, input: Documents) -> Embeddings:\n",
    "        embeddings = self(input)\n",
    "        return embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6adf5bbc5ffa44eea1166d595fe5b86a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64dbbc09bae1477b8f084adf2747c1db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at mistralai/Mistral-7B-Instruct-v0.2 were not used when initializing MistralEncoderModel: ['lm_head.weight']\n",
      "- This IS expected if you are initializing MistralEncoderModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing MistralEncoderModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "ename": "TypeError",
     "evalue": "BFloat16 is not supported on MPS",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m peftEmbeddingModel \u001b[38;5;241m=\u001b[39m \u001b[43mPeftEmbeddingModel\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[6], line 13\u001b[0m, in \u001b[0;36mPeftEmbeddingModel.__init__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m      7\u001b[0m tokenizer \u001b[38;5;241m=\u001b[39m AutoTokenizer\u001b[38;5;241m.\u001b[39mfrom_pretrained(\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMcGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      9\u001b[0m )\n\u001b[1;32m     10\u001b[0m config \u001b[38;5;241m=\u001b[39m AutoConfig\u001b[38;5;241m.\u001b[39mfrom_pretrained(\n\u001b[1;32m     11\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMcGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\u001b[39m\u001b[38;5;124m\"\u001b[39m, trust_remote_code\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m     12\u001b[0m )\n\u001b[0;32m---> 13\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mAutoModel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mMcGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     15\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtrust_remote_code\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     16\u001b[0m \u001b[43m    \u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     17\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtorch_dtype\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfloat16\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdevice_map\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmps\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mis_available\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01melse\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcpu\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     19\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     20\u001b[0m model \u001b[38;5;241m=\u001b[39m PeftModel\u001b[38;5;241m.\u001b[39mfrom_pretrained(\n\u001b[1;32m     21\u001b[0m     model,\n\u001b[1;32m     22\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMcGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m     23\u001b[0m )\n\u001b[1;32m     24\u001b[0m model \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mmerge_and_unload()  \u001b[38;5;66;03m# This can take several minutes on cpu\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:558\u001b[0m, in \u001b[0;36m_BaseAutoModelClass.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m    556\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    557\u001b[0m         \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39mregister(config\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__class__\u001b[39m, model_class, exist_ok\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m--> 558\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmodel_class\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_pretrained\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    559\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpretrained_model_name_or_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mmodel_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mhub_kwargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\n\u001b[1;32m    560\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    561\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(config) \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_model_mapping\u001b[38;5;241m.\u001b[39mkeys():\n\u001b[1;32m    562\u001b[0m     model_class \u001b[38;5;241m=\u001b[39m _get_model_class(config, \u001b[38;5;28mcls\u001b[39m\u001b[38;5;241m.\u001b[39m_model_mapping)\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/transformers/modeling_utils.py:3596\u001b[0m, in \u001b[0;36mPreTrainedModel.from_pretrained\u001b[0;34m(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, *model_args, **kwargs)\u001b[0m\n\u001b[1;32m   3593\u001b[0m     model\u001b[38;5;241m.\u001b[39mhf_quantizer \u001b[38;5;241m=\u001b[39m hf_quantizer\n\u001b[1;32m   3595\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m _adapter_model_path \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m-> 3596\u001b[0m     \u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mload_adapter\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   3597\u001b[0m \u001b[43m        \u001b[49m\u001b[43m_adapter_model_path\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3598\u001b[0m \u001b[43m        \u001b[49m\u001b[43madapter_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madapter_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3599\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtoken\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtoken\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3600\u001b[0m \u001b[43m        \u001b[49m\u001b[43madapter_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43madapter_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   3601\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   3603\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m output_loading_info:\n\u001b[1;32m   3604\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m loading_info \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/transformers/integrations/peft.py:193\u001b[0m, in \u001b[0;36mPeftAdapterMixin.load_adapter\u001b[0;34m(self, peft_model_id, adapter_name, revision, token, device_map, max_memory, offload_folder, offload_index, peft_config, adapter_state_dict, adapter_kwargs)\u001b[0m\n\u001b[1;32m    190\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_hf_peft_config_loaded \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    192\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m peft_model_id \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 193\u001b[0m     adapter_state_dict \u001b[38;5;241m=\u001b[39m \u001b[43mload_peft_weights\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpeft_model_id\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtoken\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtoken\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43madapter_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    195\u001b[0m \u001b[38;5;66;03m# We need to pre-process the state dict to remove unneeded prefixes - for backward compatibility\u001b[39;00m\n\u001b[1;32m    196\u001b[0m processed_adapter_state_dict \u001b[38;5;241m=\u001b[39m {}\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/peft/utils/save_and_load.py:326\u001b[0m, in \u001b[0;36mload_peft_weights\u001b[0;34m(model_id, device, **hf_hub_download_kwargs)\u001b[0m\n\u001b[1;32m    324\u001b[0m         adapters_weights \u001b[38;5;241m=\u001b[39m safe_load_file(filename, device\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    325\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 326\u001b[0m         adapters_weights \u001b[38;5;241m=\u001b[39m \u001b[43msafe_load_file\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilename\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    327\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    328\u001b[0m     adapters_weights \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mload(filename, map_location\u001b[38;5;241m=\u001b[39mtorch\u001b[38;5;241m.\u001b[39mdevice(device))\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/safetensors/torch.py:310\u001b[0m, in \u001b[0;36mload_file\u001b[0;34m(filename, device)\u001b[0m\n\u001b[1;32m    308\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m safe_open(filename, framework\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpt\u001b[39m\u001b[38;5;124m\"\u001b[39m, device\u001b[38;5;241m=\u001b[39mdevice) \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[1;32m    309\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m f\u001b[38;5;241m.\u001b[39mkeys():\n\u001b[0;32m--> 310\u001b[0m         result[k] \u001b[38;5;241m=\u001b[39m \u001b[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_tensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43mk\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    311\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "\u001b[0;31mTypeError\u001b[0m: BFloat16 is not supported on MPS"
     ]
    }
   ],
   "source": [
    "peftEmbeddingModel = PeftEmbeddingModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = chromadb.PersistentClient(path=\"./chromadb\")\n",
    "embedding_function = Embedding()\n",
    "\n",
    "collection = client.get_collection(\"Python-RAG\", embedding_function=embedding_function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fcc635cbf18a4e308ae67c1c5afdebc6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 8 files:   0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "llm_model, tokenizer = load(\"mlx-community/quantized-gemma-7b-it\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating questions: 100%|██████████| 121/121 [02:40<00:00,  1.33s/it]\n"
     ]
    }
   ],
   "source": [
    "data = collection.get()\n",
    "questions = []\n",
    "for index in tqdm(range(len(data[\"ids\"])), desc=\"Generating questions\"):\n",
    "    id = data[\"ids\"][index]\n",
    "    document = data[\"documents\"][index]\n",
    "    embeddings = embedding_function.encode(document)[0]\n",
    "    prompt = f\"<start_of_turn>user Create a vague short question for this text, without refering to the texts contents directly: {document} <start_of_turn>model Sure, here is the question:\"\n",
    "    question = generate(llm_model, tokenizer, prompt=prompt, verbose=False, max_tokens=1000)\n",
    "    question = question.replace(\"\\n\", \"\")\n",
    "    question = question.replace(\"*\", \"\")\n",
    "    questions.append((id, question, embeddings))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What is the purpose of the copyright notice and permission notice included in the text?\n",
      "How can pairs of data be sorted by their second element using lambda expressions?\n",
      "What is the purpose of the default argument values in the function `ask_ok`?\n",
      "What does the text following the heading \"4. More Control Flow Tools Python Tutorial, Release 3.7.0\" suggest about the content of the chapter?\n",
      "What is the purpose of the equal sign (=) in the text?\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "random_questions = random.sample(questions, 5)\n",
    "for question in random_questions:\n",
    "    print(question[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch.nn.functional import cosine_similarity\n",
    "\n",
    "class VectorStore:\n",
    "    def __init__(self, vectors, ids, device):\n",
    "        \"\"\"\n",
    "        Initializes the store with vectors and their corresponding ids.\n",
    "        \n",
    "        :param vectors: A list of numpy arrays.\n",
    "        :param ids: A list of ids corresponding to the vectors.\n",
    "        \"\"\"\n",
    "        assert len(vectors) == len(ids), \"Vectors and IDs must have the same length\"\n",
    "        \n",
    "        # Convert numpy arrays to PyTorch tensors and store them\n",
    "        self.vectors = [torch.tensor(vec, dtype=torch.float32, device=device) for vec in vectors]\n",
    "        print(len(self.vectors))\n",
    "        self.ids = ids\n",
    "\n",
    "    def find_closest_id(self, vector):\n",
    "        \"\"\"\n",
    "        Finds the ID of the stored vector closest to the given vector based on cosine similarity.\n",
    "        \n",
    "        :param vector: A numpy array representing the vector to compare.\n",
    "        :return: The id of the closest vector.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Calculate cosine similarity between the given vector and each stored vector\n",
    "        similarities = [cosine_similarity(vector.unsqueeze(0), vec.unsqueeze(0)) for vec in self.vectors]\n",
    "        \n",
    "        # Find the index of the maximum similarity\n",
    "        max_index = torch.argmax(torch.tensor(similarities)).item()\n",
    "        \n",
    "        # Return the corresponding ID\n",
    "        return self.ids[max_index], self.vectors[max_index]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "121\n"
     ]
    }
   ],
   "source": [
    "vectors = [vector for id, question, vector in questions]\n",
    "ids = [id for id, question, vector in questions]\n",
    "vector_store = VectorStore(vectors, ids, 'mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_to_device(batch, target_device: torch.device):\n",
    "    \"\"\"\n",
    "    send a pytorch batch to a device (CPU/GPU)\n",
    "    \"\"\"\n",
    "    for key in batch:\n",
    "        if isinstance(batch[key], torch.Tensor):\n",
    "            batch[key] = batch[key].to(target_device)\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 1/200: 100%|██████████| 5/5 [00:06<00:00,  1.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 1/200 loss: 0.09957718849182129 accuracy: 0.8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 2/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 2/200 loss: 0.0974433645606041 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 3/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 3/200 loss: 0.085300974547863 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 4/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 4/200 loss: 0.08502087742090225 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 5/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 5/200 loss: 0.08026508241891861 accuracy: 0.8375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 6/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 6/200 loss: 0.07770203799009323 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 7/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 7/200 loss: 0.07247059047222137 accuracy: 0.8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 8/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 8/200 loss: 0.07462989538908005 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 9/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 9/200 loss: 0.06969386339187622 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 10/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 10/200 loss: 0.06905487179756165 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 11/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 11/200 loss: 0.0710454210639 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 12/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 12/200 loss: 0.06408755481243134 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 13/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 13/200 loss: 0.06559502333402634 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 14/200: 100%|██████████| 5/5 [00:06<00:00,  1.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 14/200 loss: 0.06311161071062088 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 15/200: 100%|██████████| 5/5 [00:06<00:00,  1.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 15/200 loss: 0.06161809712648392 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 16/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 16/200 loss: 0.05626188591122627 accuracy: 0.7875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 17/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 17/200 loss: 0.06164275482296944 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 18/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 18/200 loss: 0.056734777987003326 accuracy: 0.775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 19/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 19/200 loss: 0.0583144836127758 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 20/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 20/200 loss: 0.057631779462099075 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 21/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 21/200 loss: 0.05656299740076065 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 22/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 22/200 loss: 0.054244816303253174 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 23/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 23/200 loss: 0.054403699934482574 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 24/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 24/200 loss: 0.05478975921869278 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 25/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 25/200 loss: 0.05591518431901932 accuracy: 0.6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 26/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 26/200 loss: 0.05302702262997627 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 27/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 27/200 loss: 0.05021423101425171 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 28/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 28/200 loss: 0.05388776212930679 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 29/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 29/200 loss: 0.04979405179619789 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 30/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 30/200 loss: 0.051896750926971436 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 31/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 31/200 loss: 0.05068845674395561 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 32/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 32/200 loss: 0.051299966871738434 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 33/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 33/200 loss: 0.05208313465118408 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 34/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 34/200 loss: 0.05201618745923042 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 35/200: 100%|██████████| 5/5 [00:07<00:00,  1.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 35/200 loss: 0.050849270075559616 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 36/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 36/200 loss: 0.05148733779788017 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 37/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 37/200 loss: 0.04944610595703125 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 38/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 38/200 loss: 0.050332825630903244 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 39/200: 100%|██████████| 5/5 [00:07<00:00,  1.48s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 39/200 loss: 0.04777860268950462 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 40/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 40/200 loss: 0.047355517745018005 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 41/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 41/200 loss: 0.046889133751392365 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 42/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 42/200 loss: 0.047952938824892044 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 43/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 43/200 loss: 0.04736020416021347 accuracy: 0.8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 44/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 44/200 loss: 0.050170499831438065 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 45/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 45/200 loss: 0.046396058052778244 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 46/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 46/200 loss: 0.046988438814878464 accuracy: 0.625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 47/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 47/200 loss: 0.04617663100361824 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 48/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 48/200 loss: 0.045595597475767136 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 49/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 49/200 loss: 0.04582246020436287 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 50/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 50/200 loss: 0.046335794031620026 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 51/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 51/200 loss: 0.04654260724782944 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 52/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 52/200 loss: 0.04753196984529495 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 53/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 53/200 loss: 0.0443430133163929 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 54/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 54/200 loss: 0.044266264885663986 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 55/200: 100%|██████████| 5/5 [00:07<00:00,  1.48s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 55/200 loss: 0.044693656265735626 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 56/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 56/200 loss: 0.04474736377596855 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 57/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 57/200 loss: 0.045365504920482635 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 58/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 58/200 loss: 0.04406077414751053 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 59/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 59/200 loss: 0.04267193377017975 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 60/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 60/200 loss: 0.043680980801582336 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 61/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 61/200 loss: 0.044491805136203766 accuracy: 0.7875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 62/200: 100%|██████████| 5/5 [00:07<00:00,  1.50s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 62/200 loss: 0.04269612208008766 accuracy: 0.6125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 63/200: 100%|██████████| 5/5 [00:07<00:00,  1.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 63/200 loss: 0.042902376502752304 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 64/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 64/200 loss: 0.0428871251642704 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 65/200: 100%|██████████| 5/5 [00:07<00:00,  1.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 65/200 loss: 0.04370526596903801 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 66/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 66/200 loss: 0.04363435506820679 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 67/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 67/200 loss: 0.04153662174940109 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 68/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 68/200 loss: 0.04231612756848335 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 69/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 69/200 loss: 0.04020100459456444 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 70/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 70/200 loss: 0.04011337086558342 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 71/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 71/200 loss: 0.04111240431666374 accuracy: 0.625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 72/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 72/200 loss: 0.04221606254577637 accuracy: 0.6\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 73/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 73/200 loss: 0.03951388597488403 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 74/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 74/200 loss: 0.039925530552864075 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 75/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 75/200 loss: 0.04385251924395561 accuracy: 0.8125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 76/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 76/200 loss: 0.04098017141222954 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 77/200: 100%|██████████| 5/5 [00:07<00:00,  1.49s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 77/200 loss: 0.040254175662994385 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 78/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 78/200 loss: 0.039329007267951965 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 79/200: 100%|██████████| 5/5 [00:07<00:00,  1.46s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 79/200 loss: 0.04002794250845909 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 80/200: 100%|██████████| 5/5 [00:07<00:00,  1.46s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 80/200 loss: 0.039984315633773804 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 81/200: 100%|██████████| 5/5 [00:07<00:00,  1.53s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 81/200 loss: 0.04009897634387016 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 82/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 82/200 loss: 0.038481444120407104 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 83/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 83/200 loss: 0.0404372364282608 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 84/200: 100%|██████████| 5/5 [00:07<00:00,  1.46s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 84/200 loss: 0.038621701300144196 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 85/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 85/200 loss: 0.040503401309251785 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 86/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 86/200 loss: 0.03817538544535637 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 87/200: 100%|██████████| 5/5 [00:07<00:00,  1.46s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 87/200 loss: 0.03840991482138634 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 88/200: 100%|██████████| 5/5 [00:07<00:00,  1.50s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 88/200 loss: 0.03888457641005516 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 89/200: 100%|██████████| 5/5 [00:07<00:00,  1.47s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 89/200 loss: 0.03823372721672058 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 90/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 90/200 loss: 0.039797160774469376 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 91/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 91/200 loss: 0.037610121071338654 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 92/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 92/200 loss: 0.0358041375875473 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 93/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 93/200 loss: 0.038658563047647476 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 94/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 94/200 loss: 0.03772427514195442 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 95/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 95/200 loss: 0.03697945550084114 accuracy: 0.7875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 96/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 96/200 loss: 0.03619007021188736 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 97/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 97/200 loss: 0.03736802190542221 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 98/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 98/200 loss: 0.037759896367788315 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 99/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 99/200 loss: 0.03714445233345032 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 100/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 100/200 loss: 0.03525217995047569 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 101/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 101/200 loss: 0.036549728363752365 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 102/200: 100%|██████████| 5/5 [00:07<00:00,  1.45s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 102/200 loss: 0.03540791943669319 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 103/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 103/200 loss: 0.0358341746032238 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 104/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 104/200 loss: 0.03604193404316902 accuracy: 0.775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 105/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 105/200 loss: 0.0364571213722229 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 106/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 106/200 loss: 0.03646061569452286 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 107/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 107/200 loss: 0.03438546881079674 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 108/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 108/200 loss: 0.03565068170428276 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 109/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 109/200 loss: 0.03379211947321892 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 110/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 110/200 loss: 0.03331632539629936 accuracy: 0.775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 111/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 111/200 loss: 0.03424717113375664 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 112/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 112/200 loss: 0.033522870391607285 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 113/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 113/200 loss: 0.03480082377791405 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 114/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 114/200 loss: 0.032668907195329666 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 115/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 115/200 loss: 0.03390749171376228 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 116/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 116/200 loss: 0.03391314297914505 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 117/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 117/200 loss: 0.03511761128902435 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 118/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 118/200 loss: 0.03323497995734215 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 119/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 119/200 loss: 0.034874945878982544 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 120/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 120/200 loss: 0.03436115384101868 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 121/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 121/200 loss: 0.033943288028240204 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 122/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 122/200 loss: 0.03473625332117081 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 123/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 123/200 loss: 0.031384605914354324 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 124/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 124/200 loss: 0.031777460128068924 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 125/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 125/200 loss: 0.03269451856613159 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 126/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 126/200 loss: 0.03198665753006935 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 127/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 127/200 loss: 0.032449182122945786 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 128/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 128/200 loss: 0.03297121077775955 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 129/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 129/200 loss: 0.0331985242664814 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 130/200: 100%|██████████| 5/5 [00:07<00:00,  1.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 130/200 loss: 0.02998625673353672 accuracy: 0.5875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 131/200: 100%|██████████| 5/5 [00:07<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 131/200 loss: 0.03163003548979759 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 132/200: 100%|██████████| 5/5 [00:06<00:00,  1.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 132/200 loss: 0.03289437294006348 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 133/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 133/200 loss: 0.03095608949661255 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 134/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 134/200 loss: 0.03147038817405701 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 135/200: 100%|██████████| 5/5 [00:06<00:00,  1.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 135/200 loss: 0.03153738006949425 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 136/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 136/200 loss: 0.032868217676877975 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 137/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 137/200 loss: 0.032066792249679565 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 138/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 138/200 loss: 0.03302815183997154 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 139/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 139/200 loss: 0.03155415877699852 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 140/200: 100%|██████████| 5/5 [00:06<00:00,  1.40s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 140/200 loss: 0.030845988541841507 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 141/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 141/200 loss: 0.0323631577193737 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 142/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 142/200 loss: 0.032054755836725235 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 143/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 143/200 loss: 0.031612418591976166 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 144/200: 100%|██████████| 5/5 [00:07<00:00,  1.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 144/200 loss: 0.03157773241400719 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 145/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 145/200 loss: 0.030729616060853004 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 146/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 146/200 loss: 0.029063845053315163 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 147/200: 100%|██████████| 5/5 [00:07<00:00,  1.47s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 147/200 loss: 0.030055657029151917 accuracy: 0.7875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 148/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 148/200 loss: 0.03234172984957695 accuracy: 0.8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 149/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 149/200 loss: 0.028746793046593666 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 150/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 150/200 loss: 0.029864205047488213 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 151/200: 100%|██████████| 5/5 [00:07<00:00,  1.43s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 151/200 loss: 0.03226158022880554 accuracy: 0.7625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 152/200: 100%|██████████| 5/5 [00:07<00:00,  1.44s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 152/200 loss: 0.030806338414549828 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 153/200: 100%|██████████| 5/5 [00:06<00:00,  1.39s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 153/200 loss: 0.030379969626665115 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 154/200: 100%|██████████| 5/5 [00:06<00:00,  1.34s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 154/200 loss: 0.030770665034651756 accuracy: 0.8125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 155/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 155/200 loss: 0.030031254515051842 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 156/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 156/200 loss: 0.02851925790309906 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 157/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 157/200 loss: 0.030833525583148003 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 158/200: 100%|██████████| 5/5 [00:06<00:00,  1.33s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 158/200 loss: 0.028176456689834595 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 159/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 159/200 loss: 0.02897060476243496 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 160/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 160/200 loss: 0.02871864102780819 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 161/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 161/200 loss: 0.02854323945939541 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 162/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 162/200 loss: 0.030000198632478714 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 163/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 163/200 loss: 0.029101744294166565 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 164/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 164/200 loss: 0.027536148205399513 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 165/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 165/200 loss: 0.029189346358180046 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 166/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 166/200 loss: 0.02875472418963909 accuracy: 0.775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 167/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 167/200 loss: 0.02805282175540924 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 168/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 168/200 loss: 0.027757743373513222 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 169/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 169/200 loss: 0.029542630538344383 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 170/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 170/200 loss: 0.027959490194916725 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 171/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 171/200 loss: 0.030472805723547935 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 172/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 172/200 loss: 0.0275118388235569 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 173/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 173/200 loss: 0.0295927282422781 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 174/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 174/200 loss: 0.02825266122817993 accuracy: 0.8125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 175/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 175/200 loss: 0.028327900916337967 accuracy: 0.625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 176/200: 100%|██████████| 5/5 [00:06<00:00,  1.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 176/200 loss: 0.02727987803518772 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 177/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 177/200 loss: 0.02726253867149353 accuracy: 0.775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 178/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 178/200 loss: 0.028024299070239067 accuracy: 0.7375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 179/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 179/200 loss: 0.02565539814531803 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 180/200: 100%|██████████| 5/5 [00:06<00:00,  1.36s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 180/200 loss: 0.0273252185434103 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 181/200: 100%|██████████| 5/5 [00:06<00:00,  1.33s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 181/200 loss: 0.026982281357049942 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 182/200: 100%|██████████| 5/5 [00:06<00:00,  1.28s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 182/200 loss: 0.027219945564866066 accuracy: 0.75\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 183/200: 100%|██████████| 5/5 [00:06<00:00,  1.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 183/200 loss: 0.02712741494178772 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 184/200: 100%|██████████| 5/5 [00:06<00:00,  1.29s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 184/200 loss: 0.028095562011003494 accuracy: 0.7125\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 185/200: 100%|██████████| 5/5 [00:06<00:00,  1.28s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 185/200 loss: 0.02668987773358822 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 186/200: 100%|██████████| 5/5 [00:06<00:00,  1.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 186/200 loss: 0.026452992111444473 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 187/200: 100%|██████████| 5/5 [00:06<00:00,  1.29s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 187/200 loss: 0.026751196011900902 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 188/200: 100%|██████████| 5/5 [00:06<00:00,  1.33s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 188/200 loss: 0.027423067018389702 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 189/200: 100%|██████████| 5/5 [00:06<00:00,  1.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 189/200 loss: 0.027180999517440796 accuracy: 0.65\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 190/200: 100%|██████████| 5/5 [00:06<00:00,  1.22s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 190/200 loss: 0.02710375003516674 accuracy: 0.675\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 191/200: 100%|██████████| 5/5 [00:06<00:00,  1.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 191/200 loss: 0.025478381663560867 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 192/200: 100%|██████████| 5/5 [00:06<00:00,  1.35s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 192/200 loss: 0.02512943185865879 accuracy: 0.6875\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 193/200: 100%|██████████| 5/5 [00:06<00:00,  1.32s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 193/200 loss: 0.02677147462964058 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 194/200: 100%|██████████| 5/5 [00:06<00:00,  1.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 194/200 loss: 0.027180958539247513 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 195/200: 100%|██████████| 5/5 [00:06<00:00,  1.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 195/200 loss: 0.025325072929263115 accuracy: 0.7\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 196/200: 100%|██████████| 5/5 [00:06<00:00,  1.29s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 196/200 loss: 0.026265442371368408 accuracy: 0.6375\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 197/200: 100%|██████████| 5/5 [00:06<00:00,  1.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 197/200 loss: 0.025554409250617027 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 198/200: 100%|██████████| 5/5 [00:06<00:00,  1.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 198/200 loss: 0.026861492544412613 accuracy: 0.725\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 199/200: 100%|██████████| 5/5 [00:06<00:00,  1.29s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 199/200 loss: 0.025587117299437523 accuracy: 0.6625\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Episode 200/200: 100%|██████████| 5/5 [00:06<00:00,  1.31s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Episode 200/200 loss: 0.0270228274166584 accuracy: 0.7\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import copy\n",
    "\n",
    "model_id = 'thenlper/gte-large'\n",
    "model = SentenceTransformer(model_id)\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-5, eps=1e-5)\n",
    "\n",
    "model.train()\n",
    "loss_fn = nn.CosineEmbeddingLoss()\n",
    "\n",
    "episodes = 200\n",
    "batch_size = 16\n",
    "num_batches = 5\n",
    "\n",
    "for episode in range(episodes):\n",
    "    predictions = []\n",
    "    labels = []\n",
    "    accuracies = []\n",
    "    ep_reward = 0\n",
    "\n",
    "    for _ in tqdm(range(num_batches), desc=f\"Episode {episode+1}/{episodes}\"):\n",
    "        batch = random.sample(questions, batch_size)\n",
    "        \n",
    "        for id, question, vector in batch:\n",
    "            tokenized_question = model.tokenize([question])\n",
    "            tokenized_question = batch_to_device(tokenized_question, model.device)\n",
    "            embedding = model(tokenized_question)['sentence_embedding'][0]\n",
    "            prediction, label_vector = vector_store.find_closest_id(embedding)\n",
    "            if prediction == id:\n",
    "                reward = 1\n",
    "            else:\n",
    "                reward = 0\n",
    "            predictions.append(embedding)\n",
    "            labels.append(label_vector)\n",
    "            accuracies.append(reward)\n",
    "    \n",
    "    predictions = torch.stack(predictions)\n",
    "    label_vector = torch.stack(labels)\n",
    "    optimizer.zero_grad()\n",
    "    loss = loss_fn(predictions, label_vector, torch.ones(batch_size * num_batches).to(model.device))\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    accuracy = sum(accuracies) / len(accuracies)\n",
    "    print(f\"Episode {episode+1}/{episodes} loss: {loss.item()} accuracy: {accuracy}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Expected embeddings to be a list, got [[-0.01044687 -0.02135091  0.01859619 ...  0.01973298 -0.00316574\n  -0.01486639]]",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[13], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m datapoints \u001b[38;5;241m=\u001b[39m []\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m \u001b[38;5;28mid\u001b[39m, question, vector \u001b[38;5;129;01min\u001b[39;00m questions:\n\u001b[0;32m----> 6\u001b[0m     results \u001b[38;5;241m=\u001b[39m \u001b[43mcollection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquery\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery_texts\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mquestion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_results\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_results\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m      7\u001b[0m     result_ids \u001b[38;5;241m=\u001b[39m results[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mids\u001b[39m\u001b[38;5;124m\"\u001b[39m][\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m      8\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mid\u001b[39m \u001b[38;5;129;01min\u001b[39;00m result_ids:\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/chromadb/api/models/Collection.py:327\u001b[0m, in \u001b[0;36mCollection.query\u001b[0;34m(self, query_embeddings, query_texts, query_images, query_uris, n_results, where, where_document, include)\u001b[0m\n\u001b[1;32m    325\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m valid_query_embeddings \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    326\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m query_texts \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 327\u001b[0m         valid_query_embeddings \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_embed\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvalid_query_texts\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    328\u001b[0m     \u001b[38;5;28;01melif\u001b[39;00m query_images \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    329\u001b[0m         valid_query_embeddings \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_embed(\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m=\u001b[39mvalid_query_images)\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/chromadb/api/models/Collection.py:633\u001b[0m, in \u001b[0;36mCollection._embed\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    628\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_embedding_function \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    629\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    630\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou must provide an embedding function to compute embeddings.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    631\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhttps://docs.trychroma.com/embeddings\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    632\u001b[0m     )\n\u001b[0;32m--> 633\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_embedding_function\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/chromadb/api/types.py:194\u001b[0m, in \u001b[0;36mEmbeddingFunction.__init_subclass__.<locals>.__call__\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    192\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m: EmbeddingFunction[D], \u001b[38;5;28minput\u001b[39m: D) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Embeddings:\n\u001b[1;32m    193\u001b[0m     result \u001b[38;5;241m=\u001b[39m call(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m)\n\u001b[0;32m--> 194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mvalidate_embeddings\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmaybe_cast_one_to_many_embedding\u001b[49m\u001b[43m(\u001b[49m\u001b[43mresult\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/python_rag/lib/python3.10/site-packages/chromadb/api/types.py:474\u001b[0m, in \u001b[0;36mvalidate_embeddings\u001b[0;34m(embeddings)\u001b[0m\n\u001b[1;32m    472\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Validates embeddings to ensure it is a list of list of ints, or floats\"\"\"\u001b[39;00m\n\u001b[1;32m    473\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(embeddings, \u001b[38;5;28mlist\u001b[39m):\n\u001b[0;32m--> 474\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExpected embeddings to be a list, got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00membeddings\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    475\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(embeddings) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m    476\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    477\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mExpected embeddings to be a list with at least one item, got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00membeddings\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    478\u001b[0m     )\n",
      "\u001b[0;31mValueError\u001b[0m: Expected embeddings to be a list, got [[-0.01044687 -0.02135091  0.01859619 ...  0.01973298 -0.00316574\n  -0.01486639]]"
     ]
    }
   ],
   "source": [
    "accuracies = {}\n",
    "lengths = [1, 2, 3, 4, 5]\n",
    "for n_results in lengths:\n",
    "    datapoints = []\n",
    "    for id, question, vector in questions:\n",
    "        results = collection.query(query_texts=question, n_results=n_results)\n",
    "        result_ids = results[\"ids\"][0]\n",
    "        if id in result_ids:\n",
    "            datapoints.append(1)\n",
    "        else:\n",
    "            datapoints.append(0)\n",
    "\n",
    "    accuracies[n_results] = sum(datapoints) / len(datapoints)\n",
    "accuracies"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python_rag",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}