{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "48d89cbc-3660-49b9-be37-087c1b05bf78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a776258b88794927b45bbfc8fc9bdcb0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Encoding images to latents:   0%|          | 0/60000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1f4adfe6d6794196af7755992bac151c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/143 shards):   0%|          | 0/60000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ok\n"
     ]
    }
   ],
   "source": [
    "# pip install datasets diffusers transformers\n",
    "# pip install accelerate\n",
    "# pip install flash-attn --no-build-isolation\n",
    "# git config --global credential.helper store\n",
    "# pip install -U \"huggingface_hub[cli]\"\n",
    "# huggingface-cli login\n",
    "from datasets import load_dataset, DatasetDict\n",
    "from diffusers import AutoencoderKL\n",
    "from torchvision.transforms import Resize, ToTensor, Normalize, Compose, InterpolationMode\n",
    "from transformers import AutoModel, AutoImageProcessor, AutoTokenizer\n",
    "import torch\n",
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "from pathlib import Path\n",
    "\n",
    "# ---------------- 1️⃣ Настройки ----------------\n",
    "dtype = torch.float16\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "batch_size = 64\n",
    "img_size = 64\n",
    "img_share = 0.0\n",
    "empty_share = 0.01\n",
    "\n",
    "# 1. Явно создать все необходимые директории\n",
    "cache_root = Path(\"cache\")\n",
    "(cache_root/\"datasets\").mkdir(parents=True, exist_ok=True)\n",
    "Path(\"datasets\").mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "# 2. Установить переменные среды ПЕРЕД импортом библиотек\n",
    "os.environ[\"HF_HOME\"] = str(cache_root)\n",
    "os.environ[\"HF_DATASETS_CACHE\"] = str(cache_root/\"datasets\")\n",
    "\n",
    "\n",
    "# ---------------- 2️⃣ Загрузка датасета ----------------\n",
    "dataset = load_dataset(\"mnist\", split=\"train\", cache_dir=str(cache_root))\n",
    "\n",
    "# ---------------- 3️⃣ Загрузка моделей ----------------\n",
    "vae = AutoencoderKL.from_pretrained(\"AuraDiffusion/16ch-vae\", torch_dtype=dtype).to(device).eval()\n",
    "model = AutoModel.from_pretrained(\"visheratin/mexma-siglip\", torch_dtype=dtype, trust_remote_code=True, optimized=True).to(device)\n",
    "processor = AutoImageProcessor.from_pretrained(\"visheratin/mexma-siglip\",use_fast=True)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"visheratin/mexma-siglip\")\n",
    "\n",
    "# ---------------- 4️⃣ Трансформации ----------------\n",
    "transform = Compose([\n",
    "    lambda img: img.convert(\"RGB\"), \n",
    "    Resize((img_size, img_size), interpolation=InterpolationMode.BICUBIC),  # Ресайз\n",
    "    ToTensor(),  # В тензор\n",
    "    Normalize(mean=0.5, std=0.5)  # [-1, 1]\n",
    "])\n",
    "\n",
    "# ---------------- 5️⃣ Функция обработки изображений ----------------\n",
    "def encode_images_batch(images):\n",
    "    pixel_values = torch.stack([processor(images=img, return_tensors=\"pt\")[\"pixel_values\"].squeeze(0) for img in images]).to(device, dtype)\n",
    "    \n",
    "    with torch.inference_mode():\n",
    "        image_embeddings = model.vision_model(pixel_values).pooler_output #chang on last_hidden_state  # (B, 729, 1152)\n",
    "\n",
    "    return image_embeddings.unsqueeze(1).cpu().numpy()\n",
    "\n",
    "def encode_texts_batch(texts):\n",
    "    try:\n",
    "        with torch.inference_mode():\n",
    "            text_tokenized = tokenizer(texts, return_tensors=\"pt\", padding=\"max_length\").to(device)\n",
    "            features = model.text_model(\n",
    "                input_ids=text_tokenized.input_ids, attention_mask=text_tokenized.attention_mask\n",
    "            ).last_hidden_state\n",
    "            features_proj = model.text_projector(features)\n",
    "        return features_proj.cpu().numpy()\n",
    "    except Exception as e:\n",
    "        print(f\"Ошибка при кодировании текстов: {e}\")\n",
    "        raise\n",
    "\n",
    "\n",
    "# return empty str with prob\n",
    "def maybe_empty_label(label, prob=0.01):\n",
    "    return \"\" if random.random() < prob else label\n",
    "\n",
    "\n",
    "def encode_to_latents(examples):\n",
    "    pixel_values = torch.stack([transform(img) for img in examples[\"image\"]]).to(device, dtype)  # (B, 3, 256, 256)\n",
    "    \n",
    "    # VAE Latents\n",
    "    with torch.no_grad():\n",
    "        posterior = vae.encode(pixel_values.to(device)).latent_dist.mode()\n",
    "        z = (posterior - vae.config.shift_factor) * vae.config.scaling_factor\n",
    "        latents = z.cpu().numpy()\n",
    "\n",
    "    # Преобразование числовых меток в строковые\n",
    "    text_labels = [str(lbl) for lbl in examples[\"label\"]]\n",
    "    \n",
    "    if random.random() < img_share:\n",
    "        # Image Embeddings\n",
    "        pil_images = [Image.fromarray(((img.cpu().numpy().transpose(1, 2, 0) + 1) * 127.5).astype(np.uint8)) for img in pixel_values]\n",
    "        embeddings = encode_images_batch(pil_images)\n",
    "        #print(\"image_embeddings\",embeddings.shape)\n",
    "    else:\n",
    "        text_labels_with_empty = [maybe_empty_label(lbl, empty_share) for lbl in text_labels]\n",
    "        #print(\"text_labels_with_empty\",text_labels_with_empty)\n",
    "        embeddings = encode_texts_batch(text_labels_with_empty)\n",
    "        #print(\"text_embeddings\",embeddings.shape)\n",
    "\n",
    "    return {\n",
    "        \"vae\": latents,\n",
    "        \"embeddings\": embeddings,\n",
    "        \"text\": text_labels\n",
    "    }\n",
    "\n",
    "# ---------------- 6️⃣ Обработка датасета ----------------\n",
    "limited_dataset = dataset#.select(range(10))#00000))  # Ограничиваем 1000 семплов\n",
    "encoded_dataset = limited_dataset.map(\n",
    "    encode_to_latents,\n",
    "    batched=True,\n",
    "    batch_size=batch_size,\n",
    "    remove_columns=[\"image\"],\n",
    "    desc=\"Encoding images to latents\"\n",
    ")\n",
    "\n",
    "# ---------------- 7️⃣ Сохранение ----------------\n",
    "save_path = \"datasets/mnist-te\"\n",
    "os.makedirs(save_path, exist_ok=True)\n",
    "#encoded_dataset.to_parquet(os.path.join(save_path, \"dataset.parquet\"))  # Оптимальный формат\n",
    "encoded_dataset.save_to_disk(save_path)\n",
    "print(\"ok\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c0507b7e-7dbe-43ca-999a-6dc38aa1fb40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "06e183871e9f424ea684bacfc6f948a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading dataset from disk:   0%|          | 0/143 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Форма латентного представления: (16, 8, 8)\n",
      "embedding shape: (512, 1152)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUgAAAFeCAYAAADnm4a1AAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAI0VJREFUeJzt3XmQVNX58PGnu6e7Z6ZnxVmAGZiBQTYlaEBcACFqpBLxJ7jgUlHGBXeMKTTBskB/xlIxUkWVccuiWEWsRBGRLBU1UfF1iZrlFSUqiyKLs+9bz9J93j8s5vX2cw72DCMMw/dTlarcZ07fvt1z5/FwnnvO8RljjAAAFP/hvgAAGKxIkADgQIIEAAcSJAA4kCABwIEECQAOJEgAcCBBAoADCRIAHEiQOGrt2rVLfD6frF279nBfCgYpEmQfrV27Vnw+X+//UlJSpKioSMrLy2Xfvn2H+/IG3KOPPnrYE8jhvobXX39dfD6frF+//rBdAw6PlMN9AUeqe+65R8aMGSPRaFT+8Y9/yNq1a+XNN9+Ujz76SFJTUw/35Q2YRx99VPLy8qS8vPyovgYcnUiQ/fSDH/xApk+fLiIi11xzjeTl5cmqVatk06ZNsmjRosN8dYdHW1ubRCKRw30ZwIDhn9gDZPbs2SIisnPnTk/8k08+kQsvvFCGDRsmqampMn36dNm0aZN6fWNjo/zkJz+R0tJSCYfDUlxcLFdccYXU1tb2tqmurparr75aCgsLJTU1VaZOnSpPP/205zz7x9Ueeugh+dWvfiVlZWUSDoflpJNOkvfff9/TtrKyUq688kopLi6WcDgsI0aMkPPOO0927dolIiKlpaWydetW2bx5c++Qwty5c0Xk/w81bN68WW688UYpKCiQ4uJiEREpLy+X0tJS9Rnvvvtu8fl8Kr5u3TqZMWOGpKenS25urpx++uny8ssvf+M17P/ebr31Vhk1apSEw2EZN26crFq1SuLxuPp+y8vLJTs7W3JycmTx4sXS2NioriVZ+z/Ltm3b5Ec/+pFkZ2dLfn6+rFixQowxsmfPHjnvvPMkKytLhg8fLqtXr/a8vqurS1auXCnTpk2T7OxsiUQiMnv2bHnttdfUe9XV1cnll18uWVlZvdf+wQcfWMdPk73fkBx6kANkf1LJzc3tjW3dulVmzpwpRUVFsnz5colEIvLss8/KggUL5Pnnn5eFCxeKiEhra6vMnj1bPv74Y7nqqqvku9/9rtTW1sqmTZtk7969kpeXJx0dHTJ37lzZsWOH3HzzzTJmzBh57rnnpLy8XBobG+XHP/6x53qeeeYZaWlpkeuuu058Pp88+OCDcv7558tnn30mwWBQREQuuOAC2bp1qyxdulRKS0ulurpaXnnlFdm9e7eUlpbKmjVrZOnSpZKRkSF33nmniIgUFhZ63ufGG2+U/Px8WblypbS1tfX5e/vf//1fufvuu+W0006Te+65R0KhkLz77rvy6quvytlnn33Aa2hvb5c5c+bIvn375LrrrpPRo0fL22+/LXfccYdUVFTImjVrRETEGCPnnXeevPnmm3L99dfLpEmT5IUXXpDFixf3+XoTXXzxxTJp0iR54IEH5M9//rPce++9MmzYMHniiSfkjDPOkFWrVsnvfvc7ue222+Skk06S008/XUREmpub5Te/+Y1ceumlsmTJEmlpaZHf/va3Mm/ePHnvvffkhBNOEBGReDwu5557rrz33ntyww03yMSJE+XFF1+0Xnuy9xv6wKBPnnrqKSMi5m9/+5upqakxe/bsMevXrzf5+fkmHA6bPXv29LY988wzzZQpU0w0Gu2NxeNxc9ppp5ljjz22N7Zy5UojImbDhg3q/eLxuDHGmDVr1hgRMevWrev9WVdXlzn11FNNRkaGaW5uNsYY8/nnnxsRMcccc4ypr6/vbfviiy8aETF//OMfjTHGNDQ0GBExv/jFLw74eY877jgzZ84c5/cwa9Ys09PT4/nZ4sWLTUlJiXrNXXfdZb5+y23fvt34/X6zcOFCE4vFrJ/7QNfw85//3EQiEbNt2zZPfPny5SYQCJjdu3cbY4zZuHGjERHz4IMP9rbp6ekxs2fPNiJinnrqKdfHN8YY89prrxkRMc8995z6LNdee63nnMXFxcbn85kHHnigN97Q0GDS0tLM4sWLPW07Ozs979PQ0GAKCwvNVVdd1Rt7/vnnjYiYNWvW9MZisZg544wz1LUne78hefwTu5/OOussyc/Pl1GjRsmFF14okUhENm3a1PvPzPr6enn11Vdl0aJF0tLSIrW1tVJbWyt1dXUyb9482b59e2/V+/nnn5epU6da/wu//5+kf/nLX2T48OFy6aWX9v4sGAzKLbfcIq2trbJ582bP6y6++GJPb3b/EMBnn30mIiJpaWkSCoXk9ddfl4aGhn5/D0uWLJFAINCv127cuFHi8bisXLlS/H7vrWj7p3ii5557TmbPni25ubm9329tba2cddZZEovF5I033hCRr767lJQUueGGG3pfGwgEZOnSpf267q+75pprPOecPn26GGPk6quv7o3n5OTIhAkTer/7/W1DoZCIfNVLrK+vl56eHpk+fbr8+9//7m3317/+VYLBoCxZsqQ35vf75aabbvJcR1/uNySPf2L30yOPPCLjx4+XpqYmefLJJ+WNN96QcDjc+/MdO3aIMUZWrFghK1assJ6jurpaioqKZOfOnXLBBRcc8P2++OILOfbYY1UimTRpUu/Pv2706NGe4/3Jcn8yDIfDsmrVKlm2bJkUFhbKKaecIvPnz5crrrhChg8fnsQ38JUxY8Yk3TbRzp07xe/3y+TJk/v1+u3bt8uWLVskPz/f+vPq6moR+eq7GTFihGRkZHh+PmHChH6979clfs/Z2dmSmpoqeXl5Kl5XV+eJPf3007J69Wr55JNPpLu7uzf+9e90/7Wnp6d7Xjtu3DjPcV/uNySPBNlPM2bM6K1iL1iwQGbNmiWXXXaZfPrpp5KRkdFbJLjttttk3rx51nMk3uQDydWrM1/bYePWW2+Vc889VzZu3CgvvfSSrFixQu6//3559dVX5cQTT0zqfdLS0lTM1fuLxWJJnTNZ8Xhcvv/978tPf/pT68/Hjx8/oO9nY/uek/nu161bJ+Xl5bJgwQK5/fbbpaCgQAKBgNx///2q0JeMw32/DVUkyAGw/8b+3ve+J7/85S9l+fLlMnbsWBH56p/BZ5111gFfX1ZWJh999NEB25SUlMiWLVskHo97epGffPJJ78/7o6ysTJYtWybLli2T7du3ywknnCCrV6+WdevWiUhy/9RNlJuba60QJ/Zyy8rKJB6Py3//+9/eooSN6xrKysqktbX1G7/fkpIS+fvf/y6tra2eXuSnn356wNd9m9avXy9jx46VDRs2eD7fXXfd5WlXUlIir732mrS3t3t6kTt27PC068v9huQxBjlA5s6dKzNmzJA1a9ZINBqVgoICmTt3rjzxxBNSUVGh2tfU1PT+/wsuuEA++OADeeGFF1S7/b2OH/7wh1JZWSl/+MMfen/W09MjDz/8sGRkZMicOXP6dL3t7e0SjUY9sbKyMsnMzJTOzs7eWCQS6fPjMGVlZdLU1CRbtmzpjVVUVKjPt2DBAvH7/XLPPfeox3K+3ttyXcOiRYvknXfekZdeekn9rLGxUXp6ekTkq++up6dHHnvssd6fx2Ixefjhh/v0uQbS/l7m1z/nu+++K++8846n3bx586S7u1t+/etf98bi8bg88sgjnnZ9ud+QPHqQA+j222+Xiy66SNauXSvXX3+9PPLIIzJr1iyZMmWKLFmyRMaOHStVVVXyzjvvyN69e+WDDz7ofd369evloosukquuukqmTZsm9fX1smnTJnn88cdl6tSpcu2118oTTzwh5eXl8q9//UtKS0tl/fr18tZbb8maNWskMzOzT9e6bds2OfPMM2XRokUyefJkSUlJkRdeeEGqqqrkkksu6W03bdo0eeyxx+Tee++VcePGSUFBgZxxxhkHPPcll1wiP/vZz2ThwoVyyy23SHt7uzz22GMyfvx4TwFi3Lhxcuedd8rPf/5zmT17tpx//vkSDofl/fffl5EjR8r9999/wGu4/fbbZdOmTTJ//nwpLy+XadOmSVtbm3z44Yeyfv162bVrl+Tl5cm5554rM2fOlOXLl8uuXbtk8uTJsmHDBmlqaurTdzaQ5s+fLxs2bJCFCxfKOeecI59//rk8/vjjMnnyZGltbe1tt2DBApkxY4YsW7ZMduzYIRMnTpRNmzZJfX29iHh718neb+iDw1dAPzLtf7zl/fffVz+LxWKmrKzMlJWV9T76snPnTnPFFVeY4cOHm2AwaIqKisz8+fPN+vXrPa+tq6szN998sykqKjKhUMgUFxebxYsXm9ra2t42VVVV5sorrzR5eXkmFAqZKVOmqEdU9j/mY3t8R0TMXXfdZYwxpra21tx0001m4sSJJhKJmOzsbHPyySebZ5991vOayspKc84555jMzEwjIr2P2xzoezDGmJdfftkcf/zxJhQKmQkTJph169apx3z2e/LJJ82JJ55owuGwyc3NNXPmzDGvvPLKN16DMca0tLSYO+64w4wbN86EQiGTl5dnTjvtNPPQQw+Zrq4uz/d7+eWXm6ysLJOdnW0uv/xy85///OegH/OpqanxtF28eLGJRCLqHHPmzDHHHXdc73E8Hjf33XefKSkpMeFw2Jx44onmT3/6k/URqZqaGnPZZZeZzMxMk52dbcrLy81bb71lRMT8/ve/97RN9n5DcnzGsC82cKTZuHGjLFy4UN58802ZOXPm4b6cIYsECQxyHR0dnqcFYrGYnH322fLPf/5TKisrrU8SYGAwBgkMckuXLpWOjg459dRTpbOzUzZs2CBvv/223HfffSTHbxk9SGCQe+aZZ2T16tWyY8cOiUajMm7cOLnhhhvk5ptvPtyXNuSRIAHAgecgAcCBBAkADiRIAHBIuood8OmZGkaillhcxTTbsOdADoUmO3/Y9p621/b32pI9V2K7oTgsfDC/E/Tdt30fJyPZ90v2/Mm0020iIV3pb+1sSeod6UECgAMJEgAcSJAA4ECCBACHpB8U78/CqQBw+OlatDHdlnYaPUgAcCBBAoADCRIAHEiQAODAepAAhrSAr/+zs+hBAoADCRIAHEiQAODAGCSAIS1+EHNc6EECgAMJEgAcSJAA4ECCBAAHijQAhraD2NmaHiQAOJAgAcCBBAkADiRIAHCgSANgSDMHsdc6PUgAcCBBAoADCRIAHEiQAOBAkQbAkOZjuTMAGHgkSABwIEECgAMJEgAcKNIAGNoo0gDAwCNBAoADCRIAHBiDBDDE9X8Qkh4kADiQIAHAgQQJAA4kSABwoEgDYEg7iG2x6UECgAsJEgAcSJAA4ECCBAAHijTAoKVngKSk6D/ZUGpQt/OFVKytM+o5jnV1Wt7zICoaQxA9SABwIEECgAMJEgAcSJAA4ECRBhgkAgHvn2NWJKLaTJg0QcVmzjpVxdKDmSr2f95+y3P83n/+qdq0t7R843UeefpfeKIHCQAOJEgAcCBBAoADCRIAHCjSAINEeka653hUSYlqM3P26Sp24aIFKpaWkqFi8RTvzJzP932u2uyJdujXdfeo2BGF5c4AYOCRIAHAgQQJAA4kSABwGJJFmtSwXuop3TIrwe8PJHW+9oSB6/a2Nt2IVaJwkNIjaZ7jEUUjVJtj8oapWCQ9XcV6LCuZdXZ1e44Dfv3nH7Qsp9Z5pBdp9KpxSaMHCQAOJEgAcCBBAoDDET8GaRszGT5Sj90UFY9UsXAoVcUCfj1gUVFR6Tn+dMd21aY72q1iQF9kZnrHEm1jkBnpeiy9vb1LxZrq9Ko8dXU13oBtw2jG0j3oQQKAAwkSABxIkADgQIIEAIcjv0gT0g+F51keph09ShdpsjKy9fmCeo9hv89buNm1+wvVhiINDlYkYTLDsOws3SiuqyjRjqiKdVr2vPYnvDQU1H/+B/FM9ZBEDxIAHEiQAOBAggQABxIkADgc+UUay0yarCxdfBlZWKRiw3LzVCycpos0zU1NnuOQpTAkYlnhB3Dw+XQ5JBLxzqQZdkyuapNpKywG9KpU6RlpKlZYVOh9XTis2sRts2uOcAdTeKIHCQAOJEgAcCBBAoADCRIAHI74Ik1aul6ybFSxLsiMHz9OxfIL9Oyari5dbPl02zbPsV+S26oBcAlYCiuZWZme4+JRo1SbsjH6Pi4Yma9iX+yuUrGm5lbPcXOzXhKtuyemL/aI1/8yDT1IAHAgQQKAAwkSABxIkADgMASKNHpP4LFleiB78pQpKhbw6xkxlZX7VCxxboE/wKJQODgBn+6b5OV4iy1lY8eqNrZiY3dML7VXW1OjYl/s2vWNbUzsCN8De4DRgwQABxIkADiQIAHA4Ygfg0xJ0avv5B9zjIoVFhSoWNzy34eKyr26XTzuOe6JDcWHafFtSQnqh8ILLPdjSUmJ57i4qES1SY9kqFhNTbWKVVZ8qWJVVd6Hxzs69LYM8KIHCQAOJEgAcCBBAoADCRIAHI78Io1lb9+MLL2fsG1wu6NN7yccs6xmEot7Yz52D4aDP6Dvx0nHTVKxE6eeoGJz5s7xHB+Tq/d391seMG+or1Ox+gYda231rlQV7zla9nLv/zYS9CABwIEECQAOJEgAcCBBAoDDEV+ksWwvLH5LEcW2xH0oVe8LHEuYNSMi0tPtLdKYIbh3MAZGyLJP+7hj9ao858yfp2LfOeF477nCerWpHsv9WVGpV+Vpam5WsXhPl4odHdhyAQAGHAkSABxIkADgQIIEAIcjvkgTCurlzlIssxlsg+ddUT2TIHFpMxGRUNhbzPH52Rcbdj6/7nNEIpkqVjJaL2WWm5XrDVhqgT1d+p7ds/cLFauqrlCxxJk0+Gb0IAHAgQQJAA4kSABwIEECgEMfijS2p9EP/YySzMyI53jEiBGqTc6wXBULplqKOd16D+BwaqqKqcINE2ngELDsPyMxXfiLhPV+7olLpdkmbHVE9T4ye/fqvdw//fhTFWtqaNInPCqw3BkADDgSJAA4kCABwGFQPygeTtXjNCOKiz3Ho0fpB27z8/S+2EHL/tkpIf3QbUpIjyH5fN4YQ5BwCVtWiMrKzlGxtIi+t32JN5ZlELK9o13FmhobVKymWq/w09Z+tD4ozmo+ADDgSJAA4ECCBAAHEiQAOAzqIk16epqKjSwa6TkuKRml2gyzFGnE6IFanyXW3aUfHjeWFX4Am6xMvSd7WljfxyFLzCTej5baQqflQfG2Vl24scXYKqTv6EECgAMJEgAcSJAA4ECCBACHwV2kscw2yC/I9xwXFg5XbTIyMvTJLBtod3bqAe9oNKpiXQn7CccZ7IaIpFi23igdpYuGpaN1LD1NF2kS52gZE1MtOrv03tbtnToW7T5a98AeWPQgAcCBBAkADiRIAHAgQQKAw6Au0qSl6yJNblaO5zgnJ0e1CYVDKuYP6P8WNLe0qlhTc7OKtTR728V79OA5jj5FJaNVbMYpp6rYxIkTVCycqu9RX0IhMWa5zWx7W7c161isi3t0INCDBAAHEiQAOJAgAcCBBAkADoO6SBMO6X1kMjMzPcdZWZmqTcAyw8FnWTuqublRxSorKlUscc+PWEzvZYMjmb43QineIsrwkXrG1ukzZ6nYlMnHqdjEiZNULDVs2X895p1J09ik97GuqdV7zTS36HY9Mb1s39Gh//vP2NCDBAAHEiQAOJAgAcCBBAkADn0o0tgGP7/dZb/8AV1sCYfDBzwWEQlYZs3E4/paGxr1rJmKKl2kqav3Fmm6uynSHKlycnNULHGfIxGRgsIRnuNJk3WhZdZpp6jYyTP0TJphCUv0iYiIT9+jHdEOz/HWLf9VbT788EMVq6mqUrF4z9FapLHkJMtSh8miBwkADiRIAHAgQQKAw6B+UDxgGaexxRL5LG3ilqVR2iwro9TW1qlYU8IDu4xBDgbecaVAir6Vsy0rPU2Zoh/knjL1Oyo2LDfPczxpkl6R5+STv6ti+fl6vDHFcm22obL6ulrP8ebNm1WbDz78SMWqq/TD4992feBoQQ8SABxIkADgQIIEAAcSJAA4DOoije35Tr8/Iaf79MPkYtm3uiemY92Wh2m7u3QBpj3hAV4Tj+v3HCR8fv2lJX4dQcsqSalpenuLNMvWFUHLw/uBoD5f4nX4RL8uLSuiYlkRvad5RqZesSk15J0gELJcw/DhhSo2wrIqT3GR3rc6JzfXczxqtN5e4Zi8YSqWEtB/UrbHlG21xrY27332xe7dqs1nn+3Ur2vVW4dgYNCDBAAHEiQAOJAgAcCBBAkADkkXaWwFE0stZEDZtkkwCfWRuG1peduCHpbzRyK6MOFP0YP9sUFQlPEH9HVl52apWGqqXt0oljCLKBLRRY+i4iIVKyzUs0KybAWTSJqKJRZz4jH9GxhhWUVn/LgyFRtZVKzfM+Fzpli22QiH9XcWtHyP9gJel/dcqfozmrj+TJYameMPRTfsiLZ7jmsb9KyuGstMr2hnp+X86HUQeYoeJAA4kCABwIEECQAOJEgAcEi6SPNtF2RsLLskiC+hWhSwzVxIscz2sFSZRlsKE8cdf7yK1dXVe46//PJL1aajo13F4l16iTVj+VApQe9nSLPsmZybMLNDRGRk0QgVCyTONBKRlmbv1hI5OdmqzdgxJSpWMqZUxWz7OQctM25MYjXNMpNmVImenTJp4kQVy8rWxajEIoft/uzqiKpYa8JsFRGRpjq9ZcGXFdWeY7/fUqg7Tl/r6FH6M9mW37NpTNgCpKa6VrVpatbbhLC02beHHiQAOJAgAcCBBAkADiRIAHDow3Jnh34gOBbtUjF/QrElkqFnOMTjujgSsCzTNWq0Xubqf845R8UmTvAOxldX60H9qKVI0xXVMxziMT3YH0wo0gT9+teSk60LK6WWwkp6RC8hFk+Y8ZGRqQstkTS9zFhWjp41k2JZXi5m+b67Evbt6bR8Fz09+ruoqdQzRSor9Pfd1u7dT8hWvNizd5+K7bPEdluWFatJ2Oclz1Iku/yKK1SsqFDPDvIH9e/TWIqGzQl7HzUm7McuIhKN6sITvkH/t8WmBwkALiRIAHAgQQKAw6Aeg4x26zHI5jbvWF9DfZNq01ijx26yhukxpHBQP+B87LHjVCxx6f6OqH7YuKtLX2u8R4/NxWwPige8/50KWvZRzszQY4Q5OZbVfNL0+GLcJIwb2pacsaxoE7esYtRl+Z10WsbFYgljlbF2vZVFbY0eb9y9S48H7qvSD+Z/ucfbbl+lbrN79x4Vs+0h3WAZ64vHvatEnTB1qmrT3qH3VbeNx6ZY/sxs+7S3tLZ4jju7LKv0WM6Pbw89SABwIEECgAMJEgAcSJAA4DCo98WurtUD6m+/86bnuDOq9wTeO32Gik397gkqlpWlH4QOW4ociUWTQLp+GDsWtBRHLAUZ+xYR3nZ+297TlgeL/Zb/vjU36e+jtrbRc9zarh9qb27Wxa46S/GitUU/kF1bp4stiXs1tzTq13V06GJXY0OjjjXp66irqU5oo89fZ7mu1hZdWOm2FEzSIt7fZ1aOLvKlZ+gtO3wB/TuxrTTUE9dFq8Tvo6vbcq+wcs8hRQ8SABxIkADgQIIEAAcSJAA4DOoiTVOjHpx/7x//9Bx/YZl58eWX1SpW19ioYkUjhqtYjmXGTUjNuNEFE9uCIbaVe2KWmEmYsZIS0r8WfQ0iKYFKFav4Uq98s2XLVs/xlxV6RZvaer28f2WFPn9tjW6XuKWDiEi0yzu7prtTFyXiloJDzDL7yNYucaaLsZ7LVuTQfYJIpp6lNOU7kz3Hp5x2qmpTNEKv3OOzbC1hjP5MnZaZV80t3pk0tu05cGjRgwQABxIkADiQIAHAgQQJAA6DukhjE4t5B7crK3UhYctH/1fFbEWU/PxjVCwS0QP2oZC3QOK37D1tmegixjKFQm0XLba9vvX5bSuU2WZo1DXUq9jObTs9x7aZL03NjSpWXWvZl7lRz7gxYvlQg5b+0oYN0/fBidOmeY6nTTvB8rphKubz6fP3xHSBqsqybUdtfcLvhe7LYcevAAAcSJAA4ECCBAAHEiQAOBxxRZpEtuXDdn2hZ9ckLvklIhIOBlUsxTJjxZ9QIbEVfKxb71qKKD5LgUe3sS1tpsUse8bY9jFpbfEuPWYrGnRbZp3Y2h359C8l17K3z4jhBZ7jYyz7YicW70RELPU1aWnVy8t9vPVjFduTsD93t2UvcRxa9CABwIEECQAOJEgAcDjixyBtejr1SimNnfqhZ0BEJCWkx6LTUr3bKSQ+zC8iErSsumRbVaiyUj8U/q9//1vF9uz17uPd3aPvYxxa9CABwIEECQAOJEgAcCBBAoDDkCzSAH0RsDyYn1iTsRVp/JanwnssD9xXVen93Xfu/EzF9uzxFmmi7XrfcBxa9CABwIEECQAOJEgAcCBBAoADRRrAUoBJ3C7DusKS5XW2vdCra/Q+7fu+1HuTt7V598WOWfbTxqFFDxIAHEiQAOBAggQABxIkADhQpAEsW2OYmDfot2yqYZtdYxO1bJ0Q7dCxrm7vLJx43HJhOKToQQKAAwkSABxIkADgQIIEAAeKNIDRxZZYQoHEZ1nazG/Z49z4dWElGNTtAsGA5XwJMVsNiLrNIUUPEgAcSJAA4ECCBAAHEiQAOFCkASyVj8T6iLHMavH7bEUavdxZKBhSsdycYSqWnZXpOW5raVFtYj0sgXYo0YMEAAcSJAA4kCABwIExSMC2Kk/CkKNtewXbg9w+y7hkVna2ig0fWahie3bv9hxXVlapNoxB9sNBPFxPDxIAHEiQAOBAggQABxIkADhQpAFsWy6oiK2QY2llKeakBPSfWSgYtLRLXOGHpXv6LrltMJJFDxIAHEiQAOBAggQABxIkADhQpMFRz1iKLYmxuKVN3LJVg23GjS6+iMS79ao/scR9sWO6Db6JrbDV/8INPUgAcCBBAoADCRIAHEiQAOBAkQZHPWMsxZAkJrHYhv5tRZpQWM+aSUtLU7FIesRznJKi/zxZ7uzQogcJAA4kSABwIEECgAMJEgAcKNIAlj2v9b7YSb1MApY9aTIzs1RsZHGxiuWP8O5JE/z0E9WmM9qp3xQHdDCLxtGDBAAHEiQAOJAgAcCBMUgc9fx+vdqOP2EFHlsb25YLtpWBMjMyVWx4YYGKRTLSE96T/svA6P8oJL8BAHAgQQKAAwkSABxIkADgQJEGR72AZVkef8Kj4gFLI5/lofBYT4+KhVPD+vy2wlDCFg62gg8OLXqQAOBAggQABxIkADiQIAHAgSINjnpx+94JnsNAit42weezrAJk2XLBMgdHYt3dul3CewSD+j1xaNGDBAAHEiQAOJAgAcCBBAkADhRpcNTrinaoWEtzs+e4qbFBtem2FFrClj2ww6GQiqWFUy1X4i36xGOWfR7QZ7YaXLLoQQKAAwkSABxIkADgQIIEAAeKNDjqtba1q9hnn33uOR5TUqrajJ84XsUCwRwVi/v1n1k4PV3FEvlt67ChH/r/PdKDBAAHEiQAOJAgAcCBMUgc9apralTsw60feY7zC/JUm1FlJSqWX5CvYq0trSr2+a7dKtbY0Og57umOqTbou4PZuIIeJAA4kCABwIEECQAOJEgAcPCZJDfftS0lDwwNup+QGol4jkcOH6HajC0rU7HMjDQVa7YUaSorKlSsuqrSc1xXV6faxOOs8NNnfv37NbHkCmD0IAHAgQQJAA4kSABwIEECgANFGgBDms+n+4HxOEUaADgoJEgAcCBBAoADCRIAHFjuDMCQZnz9X/CMHiQAOJAgAcCBBAkADiRIAHCgSANgiKNIAwADjgQJAA4kSABwIEECgANFGgBDms+yJ02y6EECgAMJEgAcSJAA4JD0GKRtw4X+P34JHC7Jbh3C3X149Tfj6NcZwxgkAAw4EiQAOJAgAcCBBAkADkkXaWzDowHbgGhCy7jldZZtasVYG/ZzoNbyOtupbGc6iNXZB0yy12oPfsvlNGuNI4nCR7KXkOz5E7dz7+91JdnOeq8keXqf9fq/+V2sW9ZbYkn/dvu9tf1B/fIS2P7QdULwWWLG8lpf4ndmOX0w3faeyaEHCQAOJEgAcCBBAoADCRIAHJIv0tgGjAFgCKMHCQAOJEgAcCBBAoADCRIAHEiQAOBAggQABxIkADiQIAHAgQQJAA7/D5Err3THjWEFAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 400x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔹 Текстовое описание: 4\n",
      "ok\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_from_disk\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image \n",
    "\n",
    "dtype = torch.float16\n",
    "\n",
    "# Загружаем сохраненный датасет\n",
    "loaded_dataset = load_from_disk(save_path)\n",
    "\n",
    "# Проверяем структуру датасета\n",
    "#print(\"Структура датасета:\", loaded_dataset.features)\n",
    "\n",
    "# Выбираем  пример для демонстрации\n",
    "example = loaded_dataset[2]\n",
    "\n",
    "# Выводим информацию о примере\n",
    "print(\"Форма латентного представления:\", np.array(example[\"vae\"]).shape)\n",
    "print(\"embedding shape:\", np.array(example[\"embeddings\"]).shape)\n",
    "\n",
    "# Преобразуем латентное представление в тензор PyTorch\n",
    "latent_tensor = torch.tensor(example[\"vae\"], dtype=dtype).unsqueeze(0).to(device)\n",
    "\n",
    "# Декодируем латентное представление обратно в изображение\n",
    "with torch.no_grad():\n",
    "    #reconstructed_image = vae.decode(latent_tensor).sample  # Результат — тензор\n",
    "    latent = (latent_tensor.detach() / vae.config.scaling_factor) + vae.config.shift_factor\n",
    "    reconstructed_image = vae.decode(latent).sample\n",
    "\n",
    "# Переносим тензор на CPU и преобразуем в NumPy массив\n",
    "\n",
    "reconstructed_image = reconstructed_image.squeeze(0).cpu().numpy()  # Удаляем размерность батча\n",
    "\n",
    "# Переносим каналы в правильный формат (CHW -> HWC) и нормализуем значения пикселей\n",
    "reconstructed_image = np.transpose(reconstructed_image, (1, 2, 0))\n",
    "reconstructed_image = (reconstructed_image + 1) / 2  # Нормализация в диапазон [0, 1]\n",
    "\n",
    "# Преобразуем тип данных к float32\n",
    "reconstructed_image = reconstructed_image.astype(np.float32)\n",
    "reconstructed_image = np.clip(reconstructed_image, 0.0, 1.0)\n",
    "\n",
    "# Отображаем восстановленное изображение\n",
    "plt.figure(figsize=(4, 4))\n",
    "plt.imshow(reconstructed_image)\n",
    "plt.title(f\"Reconstructed Image\")\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n",
    "print(f\"🔹 Текстовое описание: {example['text']}\")\n",
    "print(\"ok\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3818a9a7-f72c-42ea-9805-43bfd4b214a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f5336198-8925-4e03-ad81-72ee0eb5248a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HF_HOME is set to: cache\n",
      "HF_DATASETS_CACHE is set to: cache/datasets\n"
     ]
    }
   ],
   "source": [
    "# Проверьте переменные окружения\n",
    "hf_home = os.environ.get(\"HF_HOME\")\n",
    "hf_datasets_cache = os.environ.get(\"HF_DATASETS_CACHE\")\n",
    "\n",
    "if hf_home:\n",
    "    print(f\"HF_HOME is set to: {hf_home}\")\n",
    "if hf_datasets_cache:\n",
    "    print(f\"HF_DATASETS_CACHE is set to: {hf_datasets_cache}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94d7f0fe-3b27-4d08-ba35-ba7de096c635",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cache/datasets\n"
     ]
    }
   ],
   "source": [
    "print(hf_datasets_cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8206ec60-e828-4a56-b902-3b44b99b563c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting transformers\n",
      "  Downloading transformers-4.49.0-py3-none-any.whl.metadata (44 kB)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from transformers) (3.13.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.26.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.28.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (1.26.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (24.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.11/dist-packages (from transformers) (6.0.2)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers) (2024.11.6)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from transformers) (2.32.3)\n",
      "Collecting tokenizers<0.22,>=0.21 (from transformers)\n",
      "  Downloading tokenizers-0.21.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.7 kB)\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.5.2)\n",
      "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers) (4.67.1)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.26.0->transformers) (2024.2.0)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub<1.0,>=0.26.0->transformers) (4.12.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (3.10)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2.2.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests->transformers) (2024.8.30)\n",
      "Downloading transformers-4.49.0-py3-none-any.whl (10.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.0/10.0 MB\u001b[0m \u001b[31m98.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tokenizers-0.21.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.0/3.0 MB\u001b[0m \u001b[31m94.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: tokenizers, transformers\n",
      "Successfully installed tokenizers-0.21.0 transformers-4.49.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install -U transformers --break-system-packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "48e6e1d3-310c-4f67-a27b-90a54036dc4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting flash-attn\n",
      "  Downloading flash_attn-2.7.4.post1.tar.gz (6.0 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.0/6.0 MB\u001b[0m \u001b[31m94.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: torch in /usr/local/lib/python3.11/dist-packages (from flash-attn) (2.4.1+cu124)\n",
      "Collecting einops (from flash-attn)\n",
      "  Downloading einops-0.8.1-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (4.12.2)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (1.12)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (2024.2.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.99 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.4.99)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.99 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.4.99)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.99 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.4.99)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (9.1.0.70)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.4.2.65 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.4.2.65)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.2.0.44 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (11.2.0.44)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.5.119 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (10.3.5.119)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.6.0.99 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (11.6.0.99)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.3.0.142 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.3.0.142)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.20.5 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (2.20.5)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.4.99 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.4.99)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.99 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (12.4.99)\n",
      "Requirement already satisfied: triton==3.0.0 in /usr/local/lib/python3.11/dist-packages (from torch->flash-attn) (3.0.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch->flash-attn) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.11/dist-packages (from sympy->torch->flash-attn) (1.3.0)\n",
      "Downloading einops-0.8.1-py3-none-any.whl (64 kB)\n",
      "Building wheels for collected packages: flash-attn\n",
      "  Building wheel for flash-attn (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for flash-attn: filename=flash_attn-2.7.4.post1-cp311-cp311-linux_x86_64.whl size=187805408 sha256=92cf49e6f66795b6934cec0cba526ed6e45d3313de3f905d45df8773f19092a9\n",
      "  Stored in directory: /root/.cache/pip/wheels/3d/88/d8/284b89f56af7d5bf366b10d6b8e251ac8a7c7bf3f04203fb4f\n",
      "Successfully built flash-attn\n",
      "Installing collected packages: einops, flash-attn\n",
      "Successfully installed einops-0.8.1 flash-attn-2.7.4.post1\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install flash-attn --no-build-isolation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b9dd4e-789a-4162-b5b2-45c26f9b7504",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}