{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7f5e2756-dae8-4791-b69e-7dad3d0b8da8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "224911bc-facf-4cfb-87d6-86a31b737e47",
   "metadata": {},
   "source": [
    "# Proof-of-Concept: Pretraining Language Models with Predicted and Fully Frozen Token Embedding Weights\n",
    "\n",
    "This notebook and accompanying code present a systematic empirical study of large language models (LLMs) trained with **predicted and entirely frozen input embedding weights** throughout the entire training phase. \n",
    "\n",
    "The key goal of this work is to demonstrate, as a proof-of-concept, that it is possible to train powerful transformer-based language models even when the token embedding matrix is:\n",
    "- **not initialized randomly or pretrained in the standard way**\n",
    "- but is instead **pre-computed by external means** (in our case: using Unicode visual features, subword patterns, or aggregated token text features)\n",
    "- and **remains absolutely frozen during all of training and finetuning**\n",
    "\n",
    "## Experimental Setup\n",
    "\n",
    "- **Pretrain dataset:** 9 billion tokens sampled from diverse English, Russian, and Chinese corpora (including Wikipedia, SQuAD_v2, TriviaQA, HotpotQA, NQ-Open, BoolQ, CommonsenseQA, AI2 ARC).\n",
    "- **Model architectures:** Transformer models trained with frozen predicted embeddings as well as baseline models with learnable (standard) embeddings. All other hyperparameters and datasets are equivalent.\n",
    "- **SFT (\"instruction-tuning\") data:** Only ~10% of the pretraining batches included SFT-style examples. The models remain primarily pretraining LMs, not instruction-tuned models.\n",
    "- **Frozen embeddings:** The token embeddings were constructed from external sources and kept frozen for 100% of pretraining and evaluation.\n",
    "- **Unfrozen baselines:** For comparison, models were trained from identical data with fully learnable embeddings in the classic manner.\n",
    "\n",
    "## Scientific Contribution\n",
    "\n",
    "The core scientific value of these experiments is to **provide the first strong empirical demonstration** that:\n",
    "- Large transformer LMs can learn and generalize even if the input embedding space contains no semantic or distributional information and is not updated during training;\n",
    "- The higher-level semantic and abstraction capabilities emerge in the transformer layers themselves, not primarily in the learned embedding table;\n",
    "- The gap in performance between frozen, predicted-embedding models and conventional (baseline) LMs is nonzero but limited, which has ramifications for modularity, model merging/fusion, and scaling.\n",
    "\n",
    "To our knowledge, **this is the first large-scale direct comparison of fully-frozen, externally-predicted embeddings vs. standard learnable embeddings on a modern LLM pretrain task** with all other factors held equal.\n",
    "\n",
    "These models are **pretrain-only** (not strong SFT/instruction-tuned LMs). \n",
    "Downstream/benchmarked metrics such as BLEU, MMLU, ARC, etc., should be interpreted accordingly: as indications of base model potential, **not as the upper bound of possible finetuned or instruction-following performance**.\n",
    "\n",
    "For context: Performance metrics of our frozen-embedding LMs are broadly comparable to strong GPT-2-class (100M–400M) baselines under similar pretrain-only settings.\n",
    "\n",
    "## What this is **not:**\n",
    "- Not a claim for SOTA on any downstream task.\n",
    "- Not an argument that no further improvement is possible with more compute or advanced tricks.\n",
    "- Not a replacement for full SFT/finetuning in LLM pipelines.\n",
    "\n",
    "## What this **is:**  \n",
    "**A proof-of-concept and key ablation establishing that high-capacity LMs can learn to encode language and world knowledge using only externally-predicted, fully-frozen token embeddings.**\n",
    "\n",
    "**All code, models, and metrics are provided for community validation and further research.**\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8da091ae-35cf-4aef-888e-33dee366081a",
   "metadata": {},
   "outputs": [],
   "source": [
    "prefixes = ['pro_bvv_en', 'pro_bvv_unfrozen',\n",
    "            'pro_bvv_ru',\n",
    "            'pro_bvv_zh', \n",
    "            'best_bvv_ru', \n",
    "            'best_bvv_unfrozen_ru', \n",
    "            'best_bvv_zh', \n",
    "            'best_bvv_unfrozen_zh',\n",
    "            #'best_bvv_moe', \n",
    "            'max_bvv_ru', \n",
    "            'max_bvv_zh', \n",
    "            #'max_bvv_moe', \n",
    "            'nemo_bvv_ru', \n",
    "            'nemo_bvv_zh', \n",
    "            #'nemo_bvv_moe'\n",
    "           ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "58baca44-3c8b-4741-bc54-e6265797758f",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_folder = 'results/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b0db1e2e-0398-4c16-a879-d2f5593869c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "files = [f for f in os.listdir(results_folder)\n",
    "         if any(f.endswith(fix+'.txt') for fix in prefixes)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "476b3330-3d5a-48a0-a01b-2e02a8fef265",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_benchmark_file(filepath):\n",
    "    with open(filepath, 'r', encoding='utf-8') as f:\n",
    "        text = f.read()\n",
    "\n",
    "    m = re.search(r'(^\\S+).*Total parameters:\\s*([0-9.]+[MB])', text, re.MULTILINE)\n",
    "    if m:\n",
    "        model_name = f\"{m.group(1)} [{m.group(2)}]\"\n",
    "    else:\n",
    "        model_name = os.path.basename(filepath)[:-4]\n",
    "\n",
    "    regex = re.compile(\n",
    "        r\"\\s\\S+\\s([A-Z\\-]+)(?: \\[([^\\]]+)\\])?:\\s+([0-9.]+)%\", re.MULTILINE)\n",
    "    data = []\n",
    "    for m in regex.finditer(text):\n",
    "        benchmark = m.group(1)\n",
    "        topic = m.group(2)\n",
    "        percent = float(m.group(3))\n",
    "        shortname = f'{benchmark}{\" [\"+topic+\"]\" if topic else \"\"}'\n",
    "        data.append({'model': model_name, 'benchmark': shortname, 'score': percent})\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "231bb76e-9730-47fe-a5e2-8213a8cb2e3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_results = []\n",
    "for file in files:\n",
    "    filepath = os.path.join(results_folder, file)\n",
    "    all_results.extend(parse_benchmark_file(filepath))\n",
    "\n",
    "df = pd.DataFrame(all_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0d0bd993-ca24-4046-be07-39d485abd490",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_comparison(df, benchmark_keyword='MMLU'):\n",
    "    # Можно фильтровать по нужному набору бенчмарков или задач\n",
    "    subdf = df[df['benchmark'].str.startswith(benchmark_keyword)]\n",
    "    plt.figure(figsize=(12, max(5, len(subdf['benchmark'].unique())*0.35)))\n",
    "    sns.barplot(\n",
    "        data=subdf,\n",
    "        x='score', y='benchmark', hue='model',\n",
    "        orient='h'\n",
    "    )\n",
    "    plt.xlabel(\"Score (%)\")\n",
    "    plt.ylabel(\"Benchmark/Task\")\n",
    "    plt.title(f\"{benchmark_keyword} comparison\")\n",
    "    plt.legend(title=\"Model\", bbox_to_anchor=(1.05,1), loc='upper left')\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d56d6465-a7d4-421b-8902-a8753c513002",
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot_comparison(df, benchmark_keyword='BLEU')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7019c5eb-646f-4bf6-8dd3-bcc43867fdf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_heatmap(df, benchmark_startswith='MMLU'):\n",
    "    pivot = df[df['benchmark'].str.startswith(benchmark_startswith)]\n",
    "    table = pivot.pivot(index='benchmark', columns='model', values='score')\n",
    "    plt.figure(figsize=(1+1.8*len(table.columns), 1+0.42*len(table)))\n",
    "    sns.heatmap(table, annot=True, fmt=\".1f\", cmap='YlGnBu', linewidths=0.5,\n",
    "                cbar_kws=dict(label='Score (%)'), vmin=0, vmax=np.nanmax(table.values))\n",
    "    plt.title(f'{benchmark_startswith} heatmap')\n",
    "    plt.ylabel(\" \")#plt.ylabel(\"Benchmark/Task\")\n",
    "    plt.xlabel(\"Model\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "78ac7a5b-43c3-4e64-a719-126fb3d5fb7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 826752 characters]",
      "text/plain": [
       "<Figure size 2260x2536 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_heatmap(df, benchmark_startswith='MMLU')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a1075db3-124d-4c6c-80a4-ebc00afff3f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_heatmap_best_subtasks(df, score_threshold=26.0):\n",
    "    # Группируем по [benchmark, model], берём максимальный score по модели (на самом деле по файлу у тебя одна строка на пару model+subtask)\n",
    "    pivot = df.pivot(index='benchmark', columns='model', values='score')\n",
    "    # Фильтруем: выбрать те строки, где максимальное значение по строке превышает threshold\n",
    "    high_score_rows = pivot[pivot.max(axis=1) > score_threshold]\n",
    "    if high_score_rows.empty:\n",
    "        print(f\"Нет задач, где хотя бы одна модель >{score_threshold}%\")\n",
    "        return\n",
    "    plt.figure(figsize=(1+1.8*len(high_score_rows.columns), 1+0.42*len(high_score_rows)))\n",
    "    sns.heatmap(high_score_rows, annot=True, fmt=\".1f\", cmap='YlGnBu', linewidths=0.5,\n",
    "                cbar_kws=dict(label='Score (%)'), vmin=0, vmax=np.nanmax(high_score_rows.values))\n",
    "    plt.title(f\"Best results >{score_threshold}%\")\n",
    "    plt.ylabel(\" \")#plt.ylabel(\"Benchmark [Task]\")\n",
    "    plt.xlabel(\"Model\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fdbcec11-9548-4f14-91e0-15c652e3294a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 262304 characters]",
      "text/plain": [
       "<Figure size 2260x730 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_heatmap_best_subtasks(df, score_threshold=28.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0bc568c1-f90c-4c87-a105-2241e84a2c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_overview_table(df):\n",
    "    # Если нужно включить и BLEU, и любые другие \"групповые\" метрики:\n",
    "    mask = (~df['benchmark'].str.contains(r'^MMLU \\[')) | (df['benchmark'].str.startswith(\"BLEU\"))\n",
    "    overview = df[mask].copy()\n",
    "    overview = overview.sort_values(['benchmark', 'model'])\n",
    "    return overview\n",
    "    \n",
    "overview = get_overview_table(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "20844f34-34bb-46fe-b190-e857c8571062",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_overview_heatmap(overview):\n",
    "    table = overview.pivot(index='benchmark', columns='model', values='score')\n",
    "    plt.figure(figsize=(1+1.8*len(table.columns), 1+0.60*len(table)))\n",
    "    sns.heatmap(table, annot=True, fmt=\".2f\", cmap='YlGnBu', linewidths=0.5,\n",
    "                cbar_kws=dict(label='Score (%)'), vmin=0, vmax=np.nanmax(table.values))\n",
    "    plt.title(\"Benchmark by model\")\n",
    "    plt.ylabel(\" \")#\"Benchmark!\"\n",
    "    plt.xlabel(\"Model\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "df33bcdf-00ce-4afa-b224-4ad539ef536b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 138148 characters]",
      "text/plain": [
       "<Figure size 2260x520 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_overview_heatmap(overview)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5030544c-c08f-4fa8-9064-d26656bc7fa6",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
