{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "399aad70",
   "metadata": {},
   "source": [
    "# Pretain Hugging Face Dataset using Masked Learning Objective\n",
    "Look over dataset tokenization tutorial prior to running this code, which will give you some of the prerequisites you need:\n",
    "1. token_dictionary\n",
    "2. Hugging Face Tokenized Dataset\n",
    "3. Example Lengths\n",
    "\n",
    "If using one of our pretrained models, this is an uneccessary step and you will just need to tokenize your own h5ad or looom file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c074f4ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stFormer.pretrain.stFormer_pretrainer import PretrainML "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d053b4f",
   "metadata": {},
   "source": [
    "## 1.1 Create Example Lengths File"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "568ec0dd",
   "metadata": {},
   "source": [
    "This file specifies the number of tokens (genes) in each spot in your tokenized data. The maximum value should be specified by max_length in tokenization process (truncated tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f26c581",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_from_disk\n",
    "import pickle\n",
    "\n",
    "ds = load_from_disk('output/spot/visium_spot.dataset')\n",
    "\n",
    "def add_lengths(example):\n",
    "    example['length'] = len(example['input_ids'])\n",
    "    return example\n",
    "length_ds = ds.map(add_lengths,num_proc=16)\n",
    "\n",
    "lengths = length_ds['length']\n",
    "with open('output/example_lengths.pickle', 'wb') as f:\n",
    "    pickle.dump(lengths, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55fa3b23",
   "metadata": {},
   "source": [
    "## 1.2 Run Pretraining using BERT Framework"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e087cf3",
   "metadata": {},
   "source": [
    "\n",
    "1. **Masking Objective**  \n",
    "   - Randomly mask out a fraction of tokens (genes) in each sequence.  \n",
    "   - Task: Predict the original gene ID at each masked position.  \n",
    "   - Loss: Cross-entropy between predicted token distribution and true gene ID.\n",
    "\n",
    "\n",
    "   - Learns rich, unsupervised representations of spatial gene expression patterns.  \n",
    "   - Captures co-expression and neighborhood relationships without labels in neighbor mode.  \n",
    "   - Provides strong initialization for downstream tasks (e.g., cell-type classification).\n",
    "\n",
    "3. **Configuration**  \n",
    "   - A standard BERT-style architecture (hidden size, layers, heads, etc.).  \n",
    "   - Dropout, layer-norm, and positional embeddings adapted for gene sequences.  \n",
    "   - Grouped batching by sequence length for efficient GPU utilization.\n",
    "\n",
    "4. **Training Loop**  \n",
    "   - Iterates over masked sequences, computing MLM loss.  \n",
    "   - Periodic checkpointing of model weights.  \n",
    "   - Final model and tokenizer saved for later fine-tuning or inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f05d3216",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = PretrainML(\n",
    "   dataset_path=\"output/spot/visium_spot.dataset\",\n",
    "   token_dict_path=\"output/token_dictionary.pickle\",\n",
    "   example_lengths_path=\"output/example_lengths.pickle\",\n",
    "   mode='spot',\n",
    "   output_dir=\"output/pretrained_model\"\n",
    ")\n",
    "\n",
    "trainer.run_pretraining()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b60a278",
   "metadata": {},
   "source": [
    "### Run Pretraining with hyperparameter serach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b8857d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = PretrainML(\n",
    "   dataset_path=\"output/spot/visium_spot.dataset\",\n",
    "   token_dict_path=\"output/token_dictionary.pickle\",\n",
    "   example_lengths_path=\"output/example_lengths.pickle\",\n",
    "   mode='spot',\n",
    "   output_dir=\"output/pretrained_model\"\n",
    ")\n",
    "\n",
    "trainer.run_hyperparameter_train(\n",
    "    search_space={\n",
    "        \"learning_rate\":       {\"type\": \"loguniform\",  \"low\": 1e-5, \"high\": 1e-3},\n",
    "        \"per_device_train_batch_size\": {\"type\": \"categorical\", \"values\": [4, 8, 16]},\n",
    "        \"weight_decay\":        {\"type\": \"loguniform\",  \"low\": 1e-6, \"high\": 1e-2},\n",
    "    },\n",
    "    resources_per_trial={'cpu': 12,'gpu':1}\n",
    "    n_trials=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa37c60",
   "metadata": {},
   "source": [
    "## 1.3 Extract Embeddings\n",
    "This module provides helper functions and a high-level class for turning\n",
    "tokenized gene sequences into fixed-size embedding vectors using a\n",
    "pretrained transformer model.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20b7831f",
   "metadata": {},
   "source": [
    "Encapsulates the end-to-end process of:\n",
    "1. **Loading**  \n",
    "   - A pretrained model from `model_directory` (with `output_hidden_states=True`)  \n",
    "   - A HuggingFace disk‐based dataset\n",
    "   - A token dictionary (gene ↔ token ID mapping)\n",
    "2. **Batching**  \n",
    "   - Iterating in chunks of `forward_batch_size`  \n",
    "   - Extracting `input_ids` and their lengths  \n",
    "3. **Preprocessing**  \n",
    "   - Applying `pad_sequences` to each batch  \n",
    "   - Generating the `attention_mask`\n",
    "4. **Model Forward Pass**  \n",
    "   - Running the model in `eval()` mode without gradients  \n",
    "   - Gathering all hidden states\n",
    "6. **Saving**  \n",
    "   - Concatenate all batch embeddings into one tensor of shape `(N, hidden_dim)`  \n",
    "   - Write to disk as `output_prefix + \".pt\"`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "893aea4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from stFormer.tokenization.embedding_extractor import EmbeddingExtractor\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a18bfb8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "extractor = EmbeddingExtractor(\n",
    "    token_dict_path=Path('output/token_dictionary.pickle'),\n",
    "    emb_mode='cls',\n",
    "    emb_layer = -1,\n",
    "    forward_batch_size=64\n",
    "    )\n",
    "embeddins = extractor.extract_embs(\n",
    "    model_directory='output/spot/models/250422_102707_stFormer_L6_E3/final',\n",
    "    dataset_path='output/spot/visium_spot.dataset',\n",
    "    output_directory='output/spot/embeddings',\n",
    "    output_prefix='visium_spot'\n",
    "\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "stFormer",
   "language": "python",
   "name": "stformer"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}