Initial commit

.gitignore

@@ -0,0 +1,58 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+env/
+ENV/
+.venv
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+
+
+# Data files
+input.txt
+*.csv
+*.json
+
+# Jupyter Notebook
+.ipynb_checkpoints/
+*.ipynb_checkpoints/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Logs
+*.log
+logs/
+
+# Hugging Face cache
+.cache/
+hf_cache/
+

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -12,3 +12,11 @@ short_description: Smollm-135 base model trained with dummy data with speedups
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Running Locally
+
+To run the application locally, use the following command:
+
+```bash
+uv run app.py
+```

Smollm_135.ipynb

@@ -0,0 +1,1002 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "d8ad4585",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "using device: cuda\n",
+      "Model parameters: 135.18M\n",
+      "loaded 338025 tokens\n",
+      "1 epoch = 165 batches\n",
+      "Starting training...\n",
+      "step 0 | loss: 10.8725 | dt: 1274.74ms | tok/sec: 1606.61\n",
+      "step 10 | loss: 7.9941 | dt: 2626.35ms | tok/sec: 779.79\n",
+      "step 20 | loss: 6.9462 | dt: 2627.22ms | tok/sec: 779.53\n",
+      "step 30 | loss: 6.7740 | dt: 3246.48ms | tok/sec: 630.84\n",
+      "step 40 | loss: 6.9956 | dt: 2638.15ms | tok/sec: 776.30\n",
+      "step 50 | loss: 6.8735 | dt: 2615.98ms | tok/sec: 782.88\n",
+      "step 60 | loss: 6.5163 | dt: 2640.97ms | tok/sec: 775.47\n",
+      "step 70 | loss: 6.5162 | dt: 2637.36ms | tok/sec: 776.53\n",
+      "step 80 | loss: 6.4836 | dt: 2661.23ms | tok/sec: 769.57\n",
+      "step 90 | loss: 6.5255 | dt: 2643.36ms | tok/sec: 774.77\n",
+      "step 100 | loss: 6.2876 | dt: 2677.88ms | tok/sec: 764.78\n",
+      "step 110 | loss: 6.5272 | dt: 2666.03ms | tok/sec: 768.18\n",
+      "step 120 | loss: 6.1898 | dt: 2824.14ms | tok/sec: 725.18\n",
+      "step 130 | loss: 6.1057 | dt: 2646.03ms | tok/sec: 773.99\n",
+      "step 140 | loss: 5.9576 | dt: 2632.82ms | tok/sec: 777.87\n",
+      "step 150 | loss: 6.4212 | dt: 2637.88ms | tok/sec: 776.38\n",
+      "step 160 | loss: 6.4857 | dt: 2641.72ms | tok/sec: 775.25\n",
+      "step 170 | loss: 6.2381 | dt: 2624.55ms | tok/sec: 780.32\n",
+      "step 180 | loss: 5.6074 | dt: 2650.12ms | tok/sec: 772.80\n",
+      "step 190 | loss: 6.0601 | dt: 2657.61ms | tok/sec: 770.62\n",
+      "step 200 | loss: 5.5407 | dt: 2856.74ms | tok/sec: 716.90\n",
+      "step 210 | loss: 5.8250 | dt: 2647.57ms | tok/sec: 773.54\n",
+      "step 220 | loss: 6.0356 | dt: 2635.08ms | tok/sec: 777.21\n",
+      "step 230 | loss: 5.7742 | dt: 2637.31ms | tok/sec: 776.55\n",
+      "step 240 | loss: 5.8564 | dt: 2645.37ms | tok/sec: 774.18\n",
+      "step 250 | loss: 5.4802 | dt: 2660.91ms | tok/sec: 769.66\n",
+      "step 260 | loss: 5.6751 | dt: 2632.88ms | tok/sec: 777.85\n",
+      "step 270 | loss: 5.9273 | dt: 2733.56ms | tok/sec: 749.21\n",
+      "step 280 | loss: 5.9138 | dt: 2626.31ms | tok/sec: 779.80\n",
+      "step 290 | loss: 5.6861 | dt: 2638.67ms | tok/sec: 776.15\n",
+      "step 300 | loss: 5.2012 | dt: 2642.88ms | tok/sec: 774.91\n",
+      "step 310 | loss: 5.6114 | dt: 2649.90ms | tok/sec: 772.86\n",
+      "step 320 | loss: 5.0033 | dt: 2688.08ms | tok/sec: 761.88\n",
+      "step 330 | loss: 5.6259 | dt: 2682.14ms | tok/sec: 763.57\n",
+      "step 340 | loss: 5.1127 | dt: 2650.79ms | tok/sec: 772.60\n",
+      "step 350 | loss: 5.3045 | dt: 2678.00ms | tok/sec: 764.75\n",
+      "step 360 | loss: 5.2118 | dt: 2666.44ms | tok/sec: 768.06\n",
+      "step 370 | loss: 5.4723 | dt: 2639.73ms | tok/sec: 775.84\n",
+      "step 380 | loss: 5.4257 | dt: 2653.49ms | tok/sec: 771.81\n",
+      "step 390 | loss: 5.0813 | dt: 2623.92ms | tok/sec: 780.51\n",
+      "step 400 | loss: 5.0538 | dt: 2637.88ms | tok/sec: 776.38\n",
+      "step 410 | loss: 5.0351 | dt: 2708.26ms | tok/sec: 756.20\n",
+      "step 420 | loss: 5.0659 | dt: 2661.11ms | tok/sec: 769.60\n",
+      "step 430 | loss: 4.9364 | dt: 2685.40ms | tok/sec: 762.64\n",
+      "step 440 | loss: 5.3093 | dt: 2632.10ms | tok/sec: 778.08\n",
+      "step 450 | loss: 5.0675 | dt: 2653.00ms | tok/sec: 771.96\n",
+      "step 460 | loss: 4.9092 | dt: 2668.73ms | tok/sec: 767.41\n",
+      "step 470 | loss: 4.5785 | dt: 2700.87ms | tok/sec: 758.28\n",
+      "step 480 | loss: 5.1269 | dt: 2705.93ms | tok/sec: 756.86\n",
+      "step 490 | loss: 5.4856 | dt: 2653.12ms | tok/sec: 771.92\n",
+      "step 500 | loss: 5.2739 | dt: 2679.78ms | tok/sec: 764.24\n",
+      "\n",
+      "--- Generating text at step 500 ---\n",
+      "!ua any cons mocking\n",
+      "I they have we be passion so his all\n",
+      "Redpt' grip:' the grave,An thy work!\n",
+      "For embr, old Capitol, goodout than coming,\n",
+      "When thet vir Rome daughter to the chance\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 510 | loss: 4.6435 | dt: 2697.87ms | tok/sec: 759.12\n",
+      "step 520 | loss: 4.9901 | dt: 2688.89ms | tok/sec: 761.65\n",
+      "step 530 | loss: 4.5844 | dt: 2703.14ms | tok/sec: 757.64\n",
+      "step 540 | loss: 4.9779 | dt: 2663.58ms | tok/sec: 768.89\n",
+      "step 550 | loss: 5.1966 | dt: 2714.69ms | tok/sec: 754.42\n",
+      "step 560 | loss: 4.9851 | dt: 2695.66ms | tok/sec: 759.74\n",
+      "step 570 | loss: 5.2129 | dt: 2698.44ms | tok/sec: 758.96\n",
+      "step 580 | loss: 4.7397 | dt: 2714.69ms | tok/sec: 754.41\n",
+      "step 590 | loss: 4.9552 | dt: 2692.72ms | tok/sec: 760.57\n",
+      "step 600 | loss: 5.2714 | dt: 2713.71ms | tok/sec: 754.69\n",
+      "step 610 | loss: 5.2913 | dt: 2653.37ms | tok/sec: 771.85\n",
+      "step 620 | loss: 5.0200 | dt: 2701.37ms | tok/sec: 758.13\n",
+      "step 630 | loss: 4.3609 | dt: 2689.80ms | tok/sec: 761.39\n",
+      "step 640 | loss: 4.9107 | dt: 2769.70ms | tok/sec: 739.43\n",
+      "step 650 | loss: 4.3624 | dt: 2707.07ms | tok/sec: 756.54\n",
+      "step 660 | loss: 5.0022 | dt: 2720.41ms | tok/sec: 752.83\n",
+      "step 670 | loss: 4.4913 | dt: 2680.64ms | tok/sec: 764.00\n",
+      "step 680 | loss: 4.7648 | dt: 2673.28ms | tok/sec: 766.10\n",
+      "step 690 | loss: 4.6267 | dt: 2671.13ms | tok/sec: 766.72\n",
+      "step 700 | loss: 4.8468 | dt: 2683.00ms | tok/sec: 763.33\n",
+      "step 710 | loss: 4.8544 | dt: 2678.93ms | tok/sec: 764.48\n",
+      "step 720 | loss: 4.5148 | dt: 2698.22ms | tok/sec: 759.02\n",
+      "step 730 | loss: 4.4280 | dt: 2694.01ms | tok/sec: 760.21\n",
+      "step 740 | loss: 4.4265 | dt: 2681.62ms | tok/sec: 763.72\n",
+      "step 750 | loss: 4.4757 | dt: 2671.03ms | tok/sec: 766.75\n",
+      "step 760 | loss: 4.3867 | dt: 2712.16ms | tok/sec: 755.12\n",
+      "step 770 | loss: 4.8252 | dt: 2681.40ms | tok/sec: 763.78\n",
+      "step 780 | loss: 4.6916 | dt: 2861.17ms | tok/sec: 715.79\n",
+      "step 790 | loss: 4.3555 | dt: 2682.62ms | tok/sec: 763.43\n",
+      "step 800 | loss: 4.0581 | dt: 2695.10ms | tok/sec: 759.90\n",
+      "step 810 | loss: 4.5024 | dt: 2718.02ms | tok/sec: 753.49\n",
+      "step 820 | loss: 4.9491 | dt: 2688.01ms | tok/sec: 761.90\n",
+      "step 830 | loss: 4.7404 | dt: 2688.48ms | tok/sec: 761.77\n",
+      "step 840 | loss: 4.1571 | dt: 2683.11ms | tok/sec: 763.29\n",
+      "step 850 | loss: 4.2970 | dt: 2673.24ms | tok/sec: 766.11\n",
+      "step 860 | loss: 4.1351 | dt: 2673.66ms | tok/sec: 765.99\n",
+      "step 870 | loss: 4.5339 | dt: 2723.79ms | tok/sec: 751.89\n",
+      "step 880 | loss: 4.7270 | dt: 2655.06ms | tok/sec: 771.36\n",
+      "step 890 | loss: 4.5174 | dt: 2654.60ms | tok/sec: 771.49\n",
+      "step 900 | loss: 4.7254 | dt: 2671.42ms | tok/sec: 766.63\n",
+      "step 910 | loss: 4.2173 | dt: 2706.20ms | tok/sec: 756.78\n",
+      "step 920 | loss: 4.4660 | dt: 2687.33ms | tok/sec: 762.09\n",
+      "step 930 | loss: 4.8292 | dt: 2668.57ms | tok/sec: 767.45\n",
+      "step 940 | loss: 4.8714 | dt: 2681.79ms | tok/sec: 763.67\n",
+      "step 950 | loss: 4.5977 | dt: 2689.44ms | tok/sec: 761.50\n",
+      "step 960 | loss: 3.9432 | dt: 2693.87ms | tok/sec: 760.24\n",
+      "step 970 | loss: 4.4277 | dt: 2685.85ms | tok/sec: 762.51\n",
+      "step 980 | loss: 3.9359 | dt: 2665.88ms | tok/sec: 768.23\n",
+      "step 990 | loss: 4.6091 | dt: 2693.13ms | tok/sec: 760.45\n",
+      "step 1000 | loss: 4.0868 | dt: 2702.22ms | tok/sec: 757.90\n",
+      "\n",
+      "--- Generating text at step 1000 ---\n",
+      "! then he did you\n",
+      "And pray them have be devour carries meet\n",
+      "And dare man, general's true.\n",
+      "\n",
+      "LidINIUS:\n",
+      "'I can I hope with you were access with call me.\n",
+      "\n",
+      "HONAN\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 1010 | loss: 4.4185 | dt: 2691.11ms | tok/sec: 761.03\n",
+      "step 1020 | loss: 4.2212 | dt: 2687.50ms | tok/sec: 762.05\n",
+      "step 1030 | loss: 4.4616 | dt: 2670.47ms | tok/sec: 766.91\n",
+      "step 1040 | loss: 4.4004 | dt: 2662.88ms | tok/sec: 769.09\n",
+      "step 1050 | loss: 4.1682 | dt: 2705.36ms | tok/sec: 757.02\n",
+      "step 1060 | loss: 4.0242 | dt: 2722.96ms | tok/sec: 752.12\n",
+      "step 1070 | loss: 4.0509 | dt: 2717.44ms | tok/sec: 753.65\n",
+      "step 1080 | loss: 4.0544 | dt: 2667.92ms | tok/sec: 767.64\n",
+      "step 1090 | loss: 4.0126 | dt: 2740.84ms | tok/sec: 747.22\n",
+      "step 1100 | loss: 4.4749 | dt: 2762.05ms | tok/sec: 741.48\n",
+      "step 1110 | loss: 4.3578 | dt: 2634.24ms | tok/sec: 777.45\n",
+      "step 1120 | loss: 4.0779 | dt: 2677.47ms | tok/sec: 764.90\n",
+      "step 1130 | loss: 3.7411 | dt: 2685.44ms | tok/sec: 762.63\n",
+      "step 1140 | loss: 4.1268 | dt: 2690.60ms | tok/sec: 761.17\n",
+      "step 1150 | loss: 4.5661 | dt: 2719.96ms | tok/sec: 752.95\n",
+      "step 1160 | loss: 4.3289 | dt: 2666.21ms | tok/sec: 768.13\n",
+      "step 1170 | loss: 3.8129 | dt: 2671.28ms | tok/sec: 766.67\n",
+      "step 1180 | loss: 3.8706 | dt: 2717.71ms | tok/sec: 753.58\n",
+      "step 1190 | loss: 3.8226 | dt: 2669.59ms | tok/sec: 767.16\n",
+      "step 1200 | loss: 4.1762 | dt: 2666.23ms | tok/sec: 768.13\n",
+      "step 1210 | loss: 4.3757 | dt: 2708.58ms | tok/sec: 756.12\n",
+      "step 1220 | loss: 4.1351 | dt: 2685.38ms | tok/sec: 762.65\n",
+      "step 1230 | loss: 4.3564 | dt: 2687.16ms | tok/sec: 762.14\n",
+      "step 1240 | loss: 3.7988 | dt: 2675.93ms | tok/sec: 765.34\n",
+      "step 1250 | loss: 4.1403 | dt: 2686.86ms | tok/sec: 762.23\n",
+      "step 1260 | loss: 4.4226 | dt: 2667.90ms | tok/sec: 767.64\n",
+      "step 1270 | loss: 4.4887 | dt: 2739.30ms | tok/sec: 747.64\n",
+      "step 1280 | loss: 4.2127 | dt: 2691.33ms | tok/sec: 760.96\n",
+      "step 1290 | loss: 3.6873 | dt: 2672.71ms | tok/sec: 766.26\n",
+      "step 1300 | loss: 4.0008 | dt: 2696.75ms | tok/sec: 759.43\n",
+      "step 1310 | loss: 3.6128 | dt: 2662.96ms | tok/sec: 769.07\n",
+      "step 1320 | loss: 4.2900 | dt: 2656.27ms | tok/sec: 771.01\n",
+      "step 1330 | loss: 3.8012 | dt: 2660.07ms | tok/sec: 769.91\n",
+      "step 1340 | loss: 4.1523 | dt: 2642.03ms | tok/sec: 775.16\n",
+      "step 1350 | loss: 3.9300 | dt: 2683.45ms | tok/sec: 763.20\n",
+      "step 1360 | loss: 4.1542 | dt: 2639.64ms | tok/sec: 775.86\n",
+      "step 1370 | loss: 4.0666 | dt: 2660.70ms | tok/sec: 769.72\n",
+      "step 1380 | loss: 3.8962 | dt: 2655.70ms | tok/sec: 771.17\n",
+      "step 1390 | loss: 3.7736 | dt: 2637.29ms | tok/sec: 776.55\n",
+      "step 1400 | loss: 3.7997 | dt: 2647.38ms | tok/sec: 773.59\n",
+      "step 1410 | loss: 3.7525 | dt: 2629.39ms | tok/sec: 778.89\n",
+      "step 1420 | loss: 3.6936 | dt: 2688.40ms | tok/sec: 761.79\n",
+      "step 1430 | loss: 4.1674 | dt: 2638.92ms | tok/sec: 776.07\n",
+      "step 1440 | loss: 4.1046 | dt: 2648.60ms | tok/sec: 773.24\n",
+      "step 1450 | loss: 3.8131 | dt: 2636.89ms | tok/sec: 776.67\n",
+      "step 1460 | loss: 3.4866 | dt: 2659.35ms | tok/sec: 770.11\n",
+      "step 1470 | loss: 3.7860 | dt: 2638.77ms | tok/sec: 776.12\n",
+      "step 1480 | loss: 4.1924 | dt: 2632.46ms | tok/sec: 777.98\n",
+      "step 1490 | loss: 4.0126 | dt: 2635.47ms | tok/sec: 777.09\n",
+      "step 1500 | loss: 3.4828 | dt: 2633.86ms | tok/sec: 777.57\n",
+      "\n",
+      "--- Generating text at step 1500 ---\n",
+      "! we rid,\n",
+      "my thing was so done.\n",
+      "\n",
+      "GLUUS:\n",
+      "F mercy in mocked itUS;\n",
+      "But rest at Rome,iances for him: but\n",
+      "Masters are worse, proud to the people,\n",
+      "And for your\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 1510 | loss: 3.6085 | dt: 2699.83ms | tok/sec: 758.57\n",
+      "step 1520 | loss: 3.5473 | dt: 2703.50ms | tok/sec: 757.54\n",
+      "step 1530 | loss: 3.8369 | dt: 2704.69ms | tok/sec: 757.20\n",
+      "step 1540 | loss: 4.0431 | dt: 2699.75ms | tok/sec: 758.59\n",
+      "step 1550 | loss: 3.8126 | dt: 2697.98ms | tok/sec: 759.09\n",
+      "step 1560 | loss: 4.0581 | dt: 2708.62ms | tok/sec: 756.10\n",
+      "step 1570 | loss: 3.4292 | dt: 2701.01ms | tok/sec: 758.23\n",
+      "step 1580 | loss: 3.8859 | dt: 2701.13ms | tok/sec: 758.20\n",
+      "step 1590 | loss: 4.0481 | dt: 2708.01ms | tok/sec: 756.27\n",
+      "step 1600 | loss: 4.1384 | dt: 2723.99ms | tok/sec: 751.84\n",
+      "step 1610 | loss: 3.9161 | dt: 2718.54ms | tok/sec: 753.35\n",
+      "step 1620 | loss: 3.3759 | dt: 2728.75ms | tok/sec: 750.53\n",
+      "step 1630 | loss: 3.6518 | dt: 2700.32ms | tok/sec: 758.43\n",
+      "step 1640 | loss: 3.3208 | dt: 2696.71ms | tok/sec: 759.44\n",
+      "step 1650 | loss: 3.9156 | dt: 2707.93ms | tok/sec: 756.30\n",
+      "step 1660 | loss: 3.5364 | dt: 2725.21ms | tok/sec: 751.50\n",
+      "step 1670 | loss: 3.8675 | dt: 2708.14ms | tok/sec: 756.24\n",
+      "step 1680 | loss: 3.6225 | dt: 2702.13ms | tok/sec: 757.92\n",
+      "step 1690 | loss: 3.8710 | dt: 2694.93ms | tok/sec: 759.95\n",
+      "step 1700 | loss: 3.7588 | dt: 2696.90ms | tok/sec: 759.39\n",
+      "step 1710 | loss: 3.6354 | dt: 2727.15ms | tok/sec: 750.97\n",
+      "step 1720 | loss: 3.5004 | dt: 2710.80ms | tok/sec: 755.50\n",
+      "step 1730 | loss: 3.5569 | dt: 2736.50ms | tok/sec: 748.40\n",
+      "step 1740 | loss: 3.4937 | dt: 2700.68ms | tok/sec: 758.33\n",
+      "step 1750 | loss: 3.4585 | dt: 2704.23ms | tok/sec: 757.33\n",
+      "step 1760 | loss: 3.8912 | dt: 2718.99ms | tok/sec: 753.22\n",
+      "step 1770 | loss: 3.9121 | dt: 2759.54ms | tok/sec: 742.15\n",
+      "step 1780 | loss: 3.5978 | dt: 2701.65ms | tok/sec: 758.06\n",
+      "step 1790 | loss: 3.2438 | dt: 2705.15ms | tok/sec: 757.08\n",
+      "step 1800 | loss: 3.4297 | dt: 2697.91ms | tok/sec: 759.11\n",
+      "step 1810 | loss: 3.8631 | dt: 2723.87ms | tok/sec: 751.87\n",
+      "step 1820 | loss: 3.7029 | dt: 2717.25ms | tok/sec: 753.70\n",
+      "step 1830 | loss: 3.2425 | dt: 2727.89ms | tok/sec: 750.76\n",
+      "step 1840 | loss: 3.3572 | dt: 2722.87ms | tok/sec: 752.15\n",
+      "step 1850 | loss: 3.2716 | dt: 2707.57ms | tok/sec: 756.40\n",
+      "step 1860 | loss: 3.5134 | dt: 2699.27ms | tok/sec: 758.73\n",
+      "step 1870 | loss: 3.7097 | dt: 2723.70ms | tok/sec: 751.92\n",
+      "step 1880 | loss: 3.5368 | dt: 2721.65ms | tok/sec: 752.48\n",
+      "step 1890 | loss: 3.7578 | dt: 2698.47ms | tok/sec: 758.95\n",
+      "step 1900 | loss: 3.2360 | dt: 2703.24ms | tok/sec: 757.61\n",
+      "step 1910 | loss: 3.5294 | dt: 2695.89ms | tok/sec: 759.68\n",
+      "step 1920 | loss: 3.7359 | dt: 2723.54ms | tok/sec: 751.96\n",
+      "step 1930 | loss: 3.8138 | dt: 2696.43ms | tok/sec: 759.52\n",
+      "step 1940 | loss: 3.6550 | dt: 2704.04ms | tok/sec: 757.38\n",
+      "step 1950 | loss: 3.0794 | dt: 2704.37ms | tok/sec: 757.29\n",
+      "step 1960 | loss: 3.3290 | dt: 2702.23ms | tok/sec: 757.89\n",
+      "step 1970 | loss: 3.0204 | dt: 2734.40ms | tok/sec: 748.98\n",
+      "step 1980 | loss: 3.5516 | dt: 2703.91ms | tok/sec: 757.42\n",
+      "step 1990 | loss: 3.3070 | dt: 2706.01ms | tok/sec: 756.83\n",
+      "step 2000 | loss: 3.5758 | dt: 2726.00ms | tok/sec: 751.28\n",
+      "\n",
+      "--- Generating text at step 2000 ---\n",
+      "!,\n",
+      "if I put thee due by! see I do not weep\n",
+      "Thy thing in this heels hits them,\n",
+      "That neverYour valour should be in pride,\n",
+      "Call him more deity tears;Signiorio\n",
+      "In time of tears\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 2010 | loss: 3.3916 | dt: 2700.63ms | tok/sec: 758.34\n",
+      "step 2020 | loss: 3.5364 | dt: 2709.72ms | tok/sec: 755.80\n",
+      "step 2030 | loss: 3.4518 | dt: 2708.71ms | tok/sec: 756.08\n",
+      "step 2040 | loss: 3.2968 | dt: 2699.77ms | tok/sec: 758.58\n",
+      "step 2050 | loss: 3.1897 | dt: 2705.14ms | tok/sec: 757.08\n",
+      "step 2060 | loss: 3.2544 | dt: 2695.59ms | tok/sec: 759.76\n",
+      "step 2070 | loss: 3.1843 | dt: 2699.45ms | tok/sec: 758.67\n",
+      "step 2080 | loss: 3.1527 | dt: 2697.34ms | tok/sec: 759.27\n",
+      "step 2090 | loss: 3.5361 | dt: 2698.70ms | tok/sec: 758.88\n",
+      "step 2100 | loss: 3.6072 | dt: 2698.23ms | tok/sec: 759.02\n",
+      "step 2110 | loss: 3.2871 | dt: 2697.04ms | tok/sec: 759.35\n",
+      "step 2120 | loss: 2.9393 | dt: 2715.87ms | tok/sec: 754.09\n",
+      "step 2130 | loss: 3.0962 | dt: 2706.66ms | tok/sec: 756.65\n",
+      "step 2140 | loss: 3.4452 | dt: 2705.67ms | tok/sec: 756.93\n",
+      "step 2150 | loss: 3.2900 | dt: 2707.68ms | tok/sec: 756.37\n",
+      "step 2160 | loss: 2.9626 | dt: 2707.08ms | tok/sec: 756.54\n",
+      "step 2170 | loss: 3.0681 | dt: 2702.33ms | tok/sec: 757.86\n",
+      "step 2180 | loss: 2.9747 | dt: 2702.98ms | tok/sec: 757.68\n",
+      "step 2190 | loss: 3.1063 | dt: 2716.67ms | tok/sec: 753.87\n",
+      "step 2200 | loss: 3.3207 | dt: 2703.48ms | tok/sec: 757.54\n",
+      "step 2210 | loss: 3.1918 | dt: 2701.67ms | tok/sec: 758.05\n",
+      "step 2220 | loss: 3.3162 | dt: 2703.80ms | tok/sec: 757.45\n",
+      "step 2230 | loss: 2.8865 | dt: 2705.55ms | tok/sec: 756.96\n",
+      "step 2240 | loss: 3.1759 | dt: 2717.99ms | tok/sec: 753.50\n",
+      "step 2250 | loss: 3.3846 | dt: 2703.64ms | tok/sec: 757.50\n",
+      "step 2260 | loss: 3.3697 | dt: 2732.82ms | tok/sec: 749.41\n",
+      "step 2270 | loss: 3.2227 | dt: 2703.69ms | tok/sec: 757.48\n",
+      "step 2280 | loss: 2.8169 | dt: 2706.62ms | tok/sec: 756.66\n",
+      "step 2290 | loss: 2.9860 | dt: 2700.74ms | tok/sec: 758.31\n",
+      "step 2300 | loss: 2.7406 | dt: 2711.47ms | tok/sec: 755.31\n",
+      "step 2310 | loss: 3.1365 | dt: 2706.79ms | tok/sec: 756.62\n",
+      "step 2320 | loss: 2.9496 | dt: 2699.63ms | tok/sec: 758.62\n",
+      "step 2330 | loss: 3.2225 | dt: 2702.71ms | tok/sec: 757.76\n",
+      "step 2340 | loss: 3.0330 | dt: 2710.74ms | tok/sec: 755.51\n",
+      "step 2350 | loss: 3.1792 | dt: 2704.27ms | tok/sec: 757.32\n",
+      "step 2360 | loss: 3.0794 | dt: 2709.94ms | tok/sec: 755.74\n",
+      "step 2370 | loss: 2.9420 | dt: 2697.96ms | tok/sec: 759.09\n",
+      "step 2380 | loss: 2.8587 | dt: 2695.43ms | tok/sec: 759.80\n",
+      "step 2390 | loss: 2.8747 | dt: 2695.93ms | tok/sec: 759.66\n",
+      "step 2400 | loss: 2.8319 | dt: 2711.99ms | tok/sec: 755.16\n",
+      "step 2410 | loss: 2.8368 | dt: 2736.55ms | tok/sec: 748.39\n",
+      "step 2420 | loss: 3.1382 | dt: 2693.63ms | tok/sec: 760.31\n",
+      "step 2430 | loss: 3.2540 | dt: 2703.97ms | tok/sec: 757.41\n",
+      "step 2440 | loss: 2.8659 | dt: 2702.61ms | tok/sec: 757.79\n",
+      "step 2450 | loss: 2.6254 | dt: 2700.30ms | tok/sec: 758.43\n",
+      "step 2460 | loss: 2.7329 | dt: 2706.12ms | tok/sec: 756.80\n",
+      "step 2470 | loss: 3.0561 | dt: 2704.39ms | tok/sec: 757.29\n",
+      "step 2480 | loss: 2.8807 | dt: 2705.39ms | tok/sec: 757.01\n",
+      "step 2490 | loss: 2.6715 | dt: 2699.49ms | tok/sec: 758.66\n",
+      "step 2500 | loss: 2.6661 | dt: 2701.31ms | tok/sec: 758.15\n",
+      "\n",
+      "--- Generating text at step 2500 ---\n",
+      "! wasBoy\n",
+      "'s one heinous to purchase with us and fold'd\n",
+      "Than is of lips for your time his son,\n",
+      "Which my conscience did forsworn,\n",
+      "It like a friendly that owe'd their go\n",
+      "But, afterced his\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 2510 | loss: 2.6807 | dt: 2698.79ms | tok/sec: 758.86\n",
+      "step 2520 | loss: 2.8045 | dt: 2696.69ms | tok/sec: 759.45\n",
+      "step 2530 | loss: 2.9238 | dt: 2701.84ms | tok/sec: 758.00\n",
+      "step 2540 | loss: 2.7829 | dt: 2696.82ms | tok/sec: 759.41\n",
+      "step 2550 | loss: 2.8688 | dt: 2713.57ms | tok/sec: 754.72\n",
+      "step 2560 | loss: 2.5317 | dt: 2700.27ms | tok/sec: 758.44\n",
+      "step 2570 | loss: 2.8389 | dt: 2694.92ms | tok/sec: 759.95\n",
+      "step 2580 | loss: 2.8973 | dt: 2700.89ms | tok/sec: 758.27\n",
+      "step 2590 | loss: 2.9376 | dt: 2701.52ms | tok/sec: 758.09\n",
+      "step 2600 | loss: 2.8337 | dt: 2734.45ms | tok/sec: 748.96\n",
+      "step 2610 | loss: 2.3944 | dt: 2696.95ms | tok/sec: 759.38\n",
+      "step 2620 | loss: 2.5720 | dt: 2705.58ms | tok/sec: 756.95\n",
+      "step 2630 | loss: 2.3286 | dt: 2726.76ms | tok/sec: 751.07\n",
+      "step 2640 | loss: 2.6583 | dt: 2712.49ms | tok/sec: 755.03\n",
+      "step 2650 | loss: 2.5653 | dt: 2706.73ms | tok/sec: 756.63\n",
+      "step 2660 | loss: 2.8053 | dt: 2709.10ms | tok/sec: 755.97\n",
+      "step 2670 | loss: 2.6431 | dt: 2696.00ms | tok/sec: 759.64\n",
+      "step 2680 | loss: 2.6953 | dt: 2746.24ms | tok/sec: 745.75\n",
+      "step 2690 | loss: 2.6361 | dt: 2716.29ms | tok/sec: 753.97\n",
+      "step 2700 | loss: 2.4590 | dt: 2774.52ms | tok/sec: 738.15\n",
+      "step 2710 | loss: 2.3923 | dt: 2749.37ms | tok/sec: 744.90\n",
+      "step 2720 | loss: 2.4924 | dt: 2701.07ms | tok/sec: 758.22\n",
+      "step 2730 | loss: 2.4631 | dt: 2711.96ms | tok/sec: 755.17\n",
+      "step 2740 | loss: 2.4050 | dt: 2718.98ms | tok/sec: 753.22\n",
+      "step 2750 | loss: 2.6518 | dt: 2697.21ms | tok/sec: 759.30\n",
+      "step 2760 | loss: 2.7260 | dt: 2876.82ms | tok/sec: 711.90\n",
+      "step 2770 | loss: 2.3956 | dt: 2746.34ms | tok/sec: 745.72\n",
+      "step 2780 | loss: 2.1536 | dt: 2684.97ms | tok/sec: 762.76\n",
+      "step 2790 | loss: 2.2741 | dt: 2954.08ms | tok/sec: 693.28\n",
+      "step 2800 | loss: 2.5598 | dt: 2906.71ms | tok/sec: 704.58\n",
+      "step 2810 | loss: 2.2911 | dt: 2680.05ms | tok/sec: 764.16\n",
+      "step 2820 | loss: 2.1884 | dt: 2682.93ms | tok/sec: 763.34\n",
+      "step 2830 | loss: 2.2343 | dt: 2675.76ms | tok/sec: 765.39\n",
+      "step 2840 | loss: 2.2411 | dt: 2695.51ms | tok/sec: 759.78\n",
+      "step 2850 | loss: 2.2290 | dt: 2864.33ms | tok/sec: 715.00\n",
+      "step 2860 | loss: 2.2985 | dt: 2659.26ms | tok/sec: 770.14\n",
+      "step 2870 | loss: 2.2090 | dt: 2667.58ms | tok/sec: 767.74\n",
+      "step 2880 | loss: 2.2442 | dt: 2657.60ms | tok/sec: 770.62\n",
+      "step 2890 | loss: 2.0802 | dt: 2673.51ms | tok/sec: 766.03\n",
+      "step 2900 | loss: 2.2653 | dt: 2668.15ms | tok/sec: 767.57\n",
+      "step 2910 | loss: 2.2835 | dt: 2662.19ms | tok/sec: 769.29\n",
+      "step 2920 | loss: 2.1685 | dt: 2660.40ms | tok/sec: 769.81\n",
+      "step 2930 | loss: 2.1162 | dt: 2660.61ms | tok/sec: 769.75\n",
+      "step 2940 | loss: 1.8274 | dt: 2661.90ms | tok/sec: 769.38\n",
+      "step 2950 | loss: 1.8564 | dt: 2703.16ms | tok/sec: 757.63\n",
+      "step 2960 | loss: 1.7792 | dt: 2658.62ms | tok/sec: 770.32\n",
+      "step 2970 | loss: 2.0663 | dt: 2662.07ms | tok/sec: 769.33\n",
+      "step 2980 | loss: 1.9452 | dt: 2660.27ms | tok/sec: 769.85\n",
+      "step 2990 | loss: 2.1067 | dt: 2664.71ms | tok/sec: 768.56\n",
+      "step 3000 | loss: 1.9999 | dt: 2669.69ms | tok/sec: 767.13\n",
+      "\n",
+      "--- Generating text at step 3000 ---\n",
+      "! thrive lives, the rage d\n",
+      "Endiciansed in the wind O very danger!\n",
+      "\n",
+      "First Murderer:\n",
+      "What commend thee well that are this!\n",
+      "\n",
+      "First Murderer:\n",
+      "On my poor lord,\n",
+      "I'll follow thee better\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 3010 | loss: 1.9782 | dt: 2669.18ms | tok/sec: 767.28\n",
+      "step 3020 | loss: 1.9162 | dt: 2816.11ms | tok/sec: 727.24\n",
+      "step 3030 | loss: 1.7600 | dt: 2748.60ms | tok/sec: 745.11\n",
+      "step 3040 | loss: 1.7559 | dt: 3330.35ms | tok/sec: 614.95\n",
+      "step 3050 | loss: 1.8405 | dt: 2766.42ms | tok/sec: 740.31\n",
+      "step 3060 | loss: 1.8245 | dt: 2675.78ms | tok/sec: 765.38\n",
+      "step 3070 | loss: 1.7706 | dt: 2834.19ms | tok/sec: 722.61\n",
+      "step 3080 | loss: 1.8807 | dt: 2701.07ms | tok/sec: 758.22\n",
+      "step 3090 | loss: 1.9564 | dt: 2774.91ms | tok/sec: 738.04\n",
+      "step 3100 | loss: 1.6477 | dt: 2787.13ms | tok/sec: 734.81\n",
+      "step 3110 | loss: 1.5303 | dt: 2768.86ms | tok/sec: 739.65\n",
+      "step 3120 | loss: 1.5144 | dt: 2668.25ms | tok/sec: 767.54\n",
+      "step 3130 | loss: 1.6874 | dt: 2684.58ms | tok/sec: 762.87\n",
+      "step 3140 | loss: 1.5967 | dt: 2672.88ms | tok/sec: 766.22\n",
+      "step 3150 | loss: 1.4894 | dt: 2685.02ms | tok/sec: 762.75\n",
+      "step 3160 | loss: 1.5369 | dt: 2780.46ms | tok/sec: 736.57\n",
+      "step 3170 | loss: 1.5521 | dt: 2766.64ms | tok/sec: 740.25\n",
+      "step 3180 | loss: 1.5144 | dt: 2710.92ms | tok/sec: 755.46\n",
+      "step 3190 | loss: 1.5083 | dt: 2749.40ms | tok/sec: 744.89\n",
+      "step 3200 | loss: 1.4799 | dt: 2678.04ms | tok/sec: 764.74\n",
+      "step 3210 | loss: 1.5245 | dt: 2707.56ms | tok/sec: 756.40\n",
+      "step 3220 | loss: 1.3773 | dt: 2688.40ms | tok/sec: 761.79\n",
+      "step 3230 | loss: 1.5153 | dt: 2684.65ms | tok/sec: 762.86\n",
+      "step 3240 | loss: 1.4807 | dt: 2753.05ms | tok/sec: 743.90\n",
+      "step 3250 | loss: 1.4735 | dt: 2755.05ms | tok/sec: 743.36\n",
+      "step 3260 | loss: 1.3735 | dt: 2719.84ms | tok/sec: 752.98\n",
+      "step 3270 | loss: 1.2075 | dt: 2843.99ms | tok/sec: 720.12\n",
+      "step 3280 | loss: 1.2305 | dt: 2703.84ms | tok/sec: 757.44\n",
+      "step 3290 | loss: 1.1408 | dt: 2739.76ms | tok/sec: 747.51\n",
+      "step 3300 | loss: 1.2520 | dt: 2892.35ms | tok/sec: 708.07\n",
+      "step 3310 | loss: 1.3056 | dt: 2716.91ms | tok/sec: 753.80\n",
+      "step 3320 | loss: 1.2803 | dt: 2681.74ms | tok/sec: 763.68\n",
+      "step 3330 | loss: 1.2850 | dt: 2725.97ms | tok/sec: 751.29\n",
+      "step 3340 | loss: 1.2028 | dt: 2692.16ms | tok/sec: 760.73\n",
+      "step 3350 | loss: 1.1584 | dt: 2698.27ms | tok/sec: 759.00\n",
+      "step 3360 | loss: 1.0396 | dt: 2998.41ms | tok/sec: 683.03\n",
+      "step 3370 | loss: 1.0607 | dt: 2903.31ms | tok/sec: 705.40\n",
+      "step 3380 | loss: 1.1164 | dt: 3033.53ms | tok/sec: 675.12\n",
+      "step 3390 | loss: 1.1108 | dt: 2886.95ms | tok/sec: 709.40\n",
+      "step 3400 | loss: 1.0315 | dt: 2752.63ms | tok/sec: 744.02\n",
+      "step 3410 | loss: 1.1562 | dt: 2726.22ms | tok/sec: 751.22\n",
+      "step 3420 | loss: 1.1321 | dt: 2708.28ms | tok/sec: 756.20\n",
+      "step 3430 | loss: 0.9760 | dt: 2681.26ms | tok/sec: 763.82\n",
+      "step 3440 | loss: 0.8877 | dt: 2717.40ms | tok/sec: 753.66\n",
+      "step 3450 | loss: 0.9401 | dt: 2842.15ms | tok/sec: 720.58\n",
+      "step 3460 | loss: 0.9696 | dt: 2802.88ms | tok/sec: 730.68\n",
+      "step 3470 | loss: 0.8294 | dt: 2881.42ms | tok/sec: 710.76\n",
+      "step 3480 | loss: 0.9166 | dt: 2741.99ms | tok/sec: 746.90\n",
+      "step 3490 | loss: 0.8187 | dt: 2768.42ms | tok/sec: 739.77\n",
+      "step 3500 | loss: 0.8012 | dt: 2875.77ms | tok/sec: 712.16\n",
+      "\n",
+      "--- Generating text at step 3500 ---\n",
+      "! who child is Rivers; all,saving thy father knows promise wish' moan'occnard's ease that Alban comes blaced boys.\n",
+      "\n",
+      "BUCKINGHAM:\n",
+      "A greater prince, a royal cousin doth make us\n",
+      "down; and\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 3510 | loss: 0.8895 | dt: 2911.00ms | tok/sec: 703.54\n",
+      "step 3520 | loss: 0.8510 | dt: 2742.43ms | tok/sec: 746.78\n",
+      "step 3530 | loss: 0.7573 | dt: 2752.54ms | tok/sec: 744.04\n",
+      "step 3540 | loss: 0.7299 | dt: 2756.80ms | tok/sec: 742.89\n",
+      "step 3550 | loss: 0.7541 | dt: 2748.92ms | tok/sec: 745.02\n",
+      "step 3560 | loss: 0.7650 | dt: 2749.40ms | tok/sec: 744.89\n",
+      "step 3570 | loss: 0.6896 | dt: 2862.06ms | tok/sec: 715.57\n",
+      "step 3580 | loss: 0.6993 | dt: 2753.84ms | tok/sec: 743.69\n",
+      "step 3590 | loss: 0.6519 | dt: 2998.55ms | tok/sec: 683.00\n",
+      "step 3600 | loss: 0.6217 | dt: 2740.47ms | tok/sec: 747.32\n",
+      "step 3610 | loss: 0.6385 | dt: 2790.61ms | tok/sec: 733.89\n",
+      "step 3620 | loss: 0.5826 | dt: 2770.31ms | tok/sec: 739.27\n",
+      "step 3630 | loss: 0.6594 | dt: 2783.17ms | tok/sec: 735.85\n",
+      "step 3640 | loss: 0.5835 | dt: 2924.64ms | tok/sec: 700.26\n",
+      "step 3650 | loss: 0.6141 | dt: 2800.75ms | tok/sec: 731.23\n",
+      "step 3660 | loss: 0.6132 | dt: 2762.03ms | tok/sec: 741.48\n",
+      "step 3670 | loss: 0.5363 | dt: 2849.28ms | tok/sec: 718.78\n",
+      "step 3680 | loss: 0.5985 | dt: 2842.88ms | tok/sec: 720.40\n",
+      "step 3690 | loss: 0.4882 | dt: 2755.73ms | tok/sec: 743.18\n",
+      "step 3700 | loss: 0.4431 | dt: 2730.10ms | tok/sec: 750.16\n",
+      "step 3710 | loss: 0.4325 | dt: 2700.16ms | tok/sec: 758.47\n",
+      "step 3720 | loss: 0.4599 | dt: 2738.14ms | tok/sec: 747.95\n",
+      "step 3730 | loss: 0.4503 | dt: 2748.69ms | tok/sec: 745.08\n",
+      "step 3740 | loss: 0.4781 | dt: 2740.93ms | tok/sec: 747.19\n",
+      "step 3750 | loss: 0.5382 | dt: 2849.36ms | tok/sec: 718.76\n",
+      "step 3760 | loss: 0.3955 | dt: 2729.07ms | tok/sec: 750.44\n",
+      "step 3770 | loss: 0.3948 | dt: 2822.13ms | tok/sec: 725.69\n",
+      "step 3780 | loss: 0.4099 | dt: 2917.25ms | tok/sec: 702.03\n",
+      "step 3790 | loss: 0.4257 | dt: 2775.18ms | tok/sec: 737.97\n",
+      "step 3800 | loss: 0.3690 | dt: 2777.39ms | tok/sec: 737.38\n",
+      "step 3810 | loss: 0.3799 | dt: 2850.00ms | tok/sec: 718.60\n",
+      "step 3820 | loss: 0.3161 | dt: 2853.62ms | tok/sec: 717.69\n",
+      "step 3830 | loss: 0.3599 | dt: 2750.53ms | tok/sec: 744.58\n",
+      "step 3840 | loss: 0.3355 | dt: 2805.98ms | tok/sec: 729.87\n",
+      "step 3850 | loss: 0.3302 | dt: 3000.64ms | tok/sec: 682.52\n",
+      "step 3860 | loss: 0.3285 | dt: 2791.06ms | tok/sec: 733.77\n",
+      "step 3870 | loss: 0.2618 | dt: 2751.07ms | tok/sec: 744.44\n",
+      "step 3880 | loss: 0.3333 | dt: 2798.18ms | tok/sec: 731.90\n",
+      "step 3890 | loss: 0.2814 | dt: 2870.91ms | tok/sec: 713.36\n",
+      "step 3900 | loss: 0.2948 | dt: 2761.24ms | tok/sec: 741.70\n",
+      "step 3910 | loss: 0.2399 | dt: 2869.73ms | tok/sec: 713.66\n",
+      "step 3920 | loss: 0.2674 | dt: 2980.56ms | tok/sec: 687.12\n",
+      "step 3930 | loss: 0.2109 | dt: 2954.62ms | tok/sec: 693.15\n",
+      "step 3940 | loss: 0.2220 | dt: 2755.38ms | tok/sec: 743.27\n",
+      "step 3950 | loss: 0.2295 | dt: 2780.64ms | tok/sec: 736.52\n",
+      "step 3960 | loss: 0.2387 | dt: 2795.49ms | tok/sec: 732.61\n",
+      "step 3970 | loss: 0.2579 | dt: 3017.64ms | tok/sec: 678.68\n",
+      "step 3980 | loss: 0.2046 | dt: 2761.85ms | tok/sec: 741.53\n",
+      "step 3990 | loss: 0.2144 | dt: 2744.11ms | tok/sec: 746.33\n",
+      "step 4000 | loss: 0.2070 | dt: 2722.34ms | tok/sec: 752.29\n",
+      "\n",
+      "--- Generating text at step 4000 ---\n",
+      "!\n",
+      "If presently you will take the one\n",
+      "To what we will o\n",
+      "\n",
+      "'ll theretove him to known your lord.\n",
+      "\n",
+      "PRINCEE:\n",
+      "Indeed, my lord, safe still well we deserve\n",
+      "In your will.\n",
+      "\n",
+      "\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 4010 | loss: 0.2144 | dt: 2902.75ms | tok/sec: 705.54\n",
+      "step 4020 | loss: 0.2045 | dt: 2945.05ms | tok/sec: 695.40\n",
+      "step 4030 | loss: 0.1790 | dt: 3056.03ms | tok/sec: 670.15\n",
+      "step 4040 | loss: 0.1516 | dt: 2747.85ms | tok/sec: 745.31\n",
+      "step 4050 | loss: 0.1615 | dt: 2805.18ms | tok/sec: 730.08\n",
+      "step 4060 | loss: 0.1421 | dt: 2749.10ms | tok/sec: 744.97\n",
+      "step 4070 | loss: 0.1674 | dt: 2859.66ms | tok/sec: 716.17\n",
+      "step 4080 | loss: 0.1633 | dt: 2847.74ms | tok/sec: 719.17\n",
+      "step 4090 | loss: 0.1570 | dt: 2916.81ms | tok/sec: 702.14\n",
+      "step 4100 | loss: 0.1592 | dt: 3016.92ms | tok/sec: 678.84\n",
+      "step 4110 | loss: 0.1363 | dt: 2727.48ms | tok/sec: 750.88\n",
+      "step 4120 | loss: 0.1797 | dt: 2730.33ms | tok/sec: 750.09\n",
+      "step 4130 | loss: 0.1210 | dt: 2699.00ms | tok/sec: 758.80\n",
+      "step 4140 | loss: 0.1253 | dt: 2867.20ms | tok/sec: 714.29\n",
+      "step 4150 | loss: 0.1016 | dt: 2856.48ms | tok/sec: 716.97\n",
+      "step 4160 | loss: 0.1162 | dt: 2750.59ms | tok/sec: 744.57\n",
+      "step 4170 | loss: 0.1599 | dt: 2886.92ms | tok/sec: 709.41\n",
+      "step 4180 | loss: 0.1116 | dt: 2701.22ms | tok/sec: 758.18\n",
+      "step 4190 | loss: 0.1093 | dt: 2674.94ms | tok/sec: 765.63\n",
+      "step 4200 | loss: 0.1099 | dt: 2875.45ms | tok/sec: 712.24\n",
+      "step 4210 | loss: 0.1319 | dt: 2775.36ms | tok/sec: 737.92\n",
+      "step 4220 | loss: 0.0932 | dt: 2740.37ms | tok/sec: 747.34\n",
+      "step 4230 | loss: 0.1047 | dt: 2697.50ms | tok/sec: 759.22\n",
+      "step 4240 | loss: 0.1159 | dt: 2656.54ms | tok/sec: 770.93\n",
+      "step 4250 | loss: 0.0943 | dt: 2632.27ms | tok/sec: 778.04\n",
+      "step 4260 | loss: 0.0935 | dt: 2733.61ms | tok/sec: 749.19\n",
+      "step 4270 | loss: 0.1044 | dt: 2700.76ms | tok/sec: 758.31\n",
+      "step 4280 | loss: 0.0909 | dt: 2690.66ms | tok/sec: 761.15\n",
+      "step 4290 | loss: 0.0993 | dt: 2689.09ms | tok/sec: 761.60\n",
+      "step 4300 | loss: 0.0931 | dt: 2710.26ms | tok/sec: 755.65\n",
+      "step 4310 | loss: 0.0841 | dt: 2795.26ms | tok/sec: 732.67\n",
+      "step 4320 | loss: 0.0729 | dt: 2704.03ms | tok/sec: 757.39\n",
+      "step 4330 | loss: 0.0633 | dt: 2676.90ms | tok/sec: 765.07\n",
+      "step 4340 | loss: 0.0859 | dt: 2761.68ms | tok/sec: 741.58\n",
+      "step 4350 | loss: 0.0625 | dt: 2778.00ms | tok/sec: 737.22\n",
+      "step 4360 | loss: 0.0638 | dt: 3028.06ms | tok/sec: 676.34\n",
+      "step 4370 | loss: 0.0841 | dt: 2895.60ms | tok/sec: 707.28\n",
+      "step 4380 | loss: 0.0731 | dt: 2678.41ms | tok/sec: 764.63\n",
+      "step 4390 | loss: 0.0676 | dt: 2632.73ms | tok/sec: 777.90\n",
+      "step 4400 | loss: 0.0789 | dt: 2633.57ms | tok/sec: 777.65\n",
+      "step 4410 | loss: 0.0870 | dt: 2633.24ms | tok/sec: 777.75\n",
+      "step 4420 | loss: 0.0564 | dt: 2666.58ms | tok/sec: 768.03\n",
+      "step 4430 | loss: 0.0510 | dt: 2663.78ms | tok/sec: 768.83\n",
+      "step 4440 | loss: 0.0844 | dt: 2637.76ms | tok/sec: 776.42\n",
+      "step 4450 | loss: 0.0574 | dt: 2649.37ms | tok/sec: 773.02\n",
+      "step 4460 | loss: 0.0670 | dt: 2642.94ms | tok/sec: 774.89\n",
+      "step 4470 | loss: 0.0818 | dt: 2635.25ms | tok/sec: 777.15\n",
+      "step 4480 | loss: 0.0611 | dt: 2641.49ms | tok/sec: 775.32\n",
+      "step 4490 | loss: 0.0571 | dt: 2718.30ms | tok/sec: 753.41\n",
+      "step 4500 | loss: 0.0524 | dt: 2714.53ms | tok/sec: 754.46\n",
+      "\n",
+      "--- Generating text at step 4500 ---\n",
+      "! to Tyenting in their hate.\n",
+      "\n",
+      "BUCKINGHAM:\n",
+      "My lord, I'll bear unto your head\n",
+      "Which best lives hath had but your breath or\n",
+      "Plured with the hawnes with yields that\n",
+      "There did touch his\n",
+      "-----------------------------------\n",
+      "\n",
+      "step 4510 | loss: 0.0513 | dt: 2759.85ms | tok/sec: 742.07\n",
+      "step 4520 | loss: 0.0582 | dt: 2780.63ms | tok/sec: 736.52\n",
+      "step 4530 | loss: 0.0607 | dt: 2788.91ms | tok/sec: 734.34\n",
+      "step 4540 | loss: 0.0537 | dt: 2705.52ms | tok/sec: 756.97\n",
+      "step 4550 | loss: 0.0469 | dt: 2797.45ms | tok/sec: 732.10\n",
+      "step 4560 | loss: 0.0513 | dt: 2715.12ms | tok/sec: 754.29\n",
+      "step 4570 | loss: 0.0537 | dt: 2701.81ms | tok/sec: 758.01\n",
+      "step 4580 | loss: 0.0555 | dt: 2703.04ms | tok/sec: 757.66\n",
+      "step 4590 | loss: 0.0617 | dt: 2702.77ms | tok/sec: 757.74\n",
+      "step 4600 | loss: 0.0601 | dt: 2703.01ms | tok/sec: 757.67\n",
+      "step 4610 | loss: 0.0386 | dt: 2713.98ms | tok/sec: 754.61\n",
+      "step 4620 | loss: 0.0469 | dt: 2702.42ms | tok/sec: 757.84\n",
+      "step 4630 | loss: 0.0429 | dt: 2701.25ms | tok/sec: 758.17\n",
+      "step 4640 | loss: 0.0436 | dt: 2700.08ms | tok/sec: 758.50\n",
+      "step 4650 | loss: 0.0552 | dt: 2707.12ms | tok/sec: 756.52\n",
+      "step 4660 | loss: 0.0478 | dt: 2704.64ms | tok/sec: 757.22\n",
+      "step 4670 | loss: 0.0503 | dt: 2700.75ms | tok/sec: 758.31\n",
+      "step 4680 | loss: 0.0370 | dt: 2703.38ms | tok/sec: 757.57\n",
+      "step 4690 | loss: 0.0488 | dt: 2690.58ms | tok/sec: 761.18\n",
+      "step 4700 | loss: 0.0395 | dt: 2711.10ms | tok/sec: 755.41\n",
+      "step 4710 | loss: 0.0384 | dt: 2695.39ms | tok/sec: 759.82\n",
+      "step 4720 | loss: 0.0320 | dt: 2698.22ms | tok/sec: 759.02\n",
+      "step 4730 | loss: 0.0389 | dt: 2693.92ms | tok/sec: 760.23\n",
+      "step 4740 | loss: 0.0417 | dt: 2715.99ms | tok/sec: 754.05\n",
+      "step 4750 | loss: 0.0327 | dt: 2707.00ms | tok/sec: 756.56\n",
+      "step 4760 | loss: 0.0383 | dt: 2702.06ms | tok/sec: 757.94\n",
+      "step 4770 | loss: 0.0482 | dt: 2720.02ms | tok/sec: 752.93\n",
+      "step 4780 | loss: 0.0285 | dt: 2695.32ms | tok/sec: 759.84\n",
+      "step 4790 | loss: 0.0466 | dt: 2699.53ms | tok/sec: 758.65\n",
+      "step 4800 | loss: 0.0293 | dt: 2747.82ms | tok/sec: 745.32\n",
+      "step 4810 | loss: 0.0475 | dt: 2694.82ms | tok/sec: 759.98\n",
+      "step 4820 | loss: 0.0331 | dt: 2697.06ms | tok/sec: 759.34\n",
+      "step 4830 | loss: 0.0424 | dt: 2695.18ms | tok/sec: 759.88\n",
+      "step 4840 | loss: 0.0388 | dt: 2699.61ms | tok/sec: 758.63\n",
+      "step 4850 | loss: 0.0278 | dt: 2697.17ms | tok/sec: 759.31\n",
+      "step 4860 | loss: 0.0352 | dt: 2712.98ms | tok/sec: 754.89\n",
+      "step 4870 | loss: 0.0231 | dt: 2708.67ms | tok/sec: 756.09\n",
+      "step 4880 | loss: 0.0355 | dt: 2697.11ms | tok/sec: 759.33\n",
+      "step 4890 | loss: 0.0379 | dt: 2696.41ms | tok/sec: 759.53\n",
+      "step 4900 | loss: 0.0251 | dt: 2692.32ms | tok/sec: 760.68\n",
+      "step 4910 | loss: 0.0263 | dt: 2695.45ms | tok/sec: 759.80\n",
+      "step 4920 | loss: 0.0279 | dt: 2703.49ms | tok/sec: 757.54\n",
+      "step 4930 | loss: 0.0253 | dt: 2696.46ms | tok/sec: 759.51\n",
+      "step 4940 | loss: 0.0279 | dt: 2719.05ms | tok/sec: 753.21\n",
+      "step 4950 | loss: 0.0303 | dt: 2698.82ms | tok/sec: 758.85\n",
+      "step 4960 | loss: 0.0277 | dt: 2702.86ms | tok/sec: 757.72\n",
+      "step 4970 | loss: 0.0230 | dt: 2708.92ms | tok/sec: 756.02\n",
+      "step 4980 | loss: 0.0397 | dt: 2693.32ms | tok/sec: 760.40\n",
+      "step 4990 | loss: 0.0244 | dt: 2699.65ms | tok/sec: 758.62\n",
+      "Saving model to smollm_135_checkpoint.pth\n",
+      "\n",
+      "--- Resuming training from checkpoint ---\n",
+      "Checkpoint loaded successfully.\n",
+      "Resume step 0 | loss: 0.0218\n",
+      "Resume step 10 | loss: 0.0335\n",
+      "Resume step 20 | loss: 0.0360\n",
+      "Resume step 30 | loss: 0.0450\n",
+      "Resume step 40 | loss: 0.0396\n",
+      "Resumed training completed.\n"
+     ]
+    }
+   ],
+   "source": [
+    "# SmolLM-135M Implementation (Llama Architecture)\n",
+    "# Based on: https://huggingface.co/HuggingFaceTB/SmolLM-135M\n",
+    "\n",
+    "import math\n",
+    "import inspect\n",
+    "from dataclasses import dataclass\n",
+    "from typing import Optional, Tuple\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.nn import functional as F\n",
+    "\n",
+    "# Configuration for SmolLM-135M\n",
+    "@dataclass\n",
+    "class SmolLMConfig:\n",
+    "    block_size: int = 512 # Reduced to 512 for 4GB GPU training\n",
+    "    vocab_size: int = 50304 # Aligned to 50304 for tiktoken compatibility (SmolLM native is 49152)\n",
+    "    n_layer: int = 30\n",
+    "    n_head: int = 9\n",
+    "    n_kv_head: int = 3 # Grouped Query Attention (GQA)\n",
+    "    n_embd: int = 576\n",
+    "    intermediate_size: int = 1536 # SwiGLU intermediate size\n",
+    "    rms_norm_eps: float = 1e-5\n",
+    "    rope_theta: float = 10000.0\n",
+    "    dropout: float = 0.0\n",
+    "    bias: bool = False # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster\n",
+    "\n",
+    "class RMSNorm(nn.Module):\n",
+    "    def __init__(self, dim: int, eps: float = 1e-6):\n",
+    "        super().__init__()\n",
+    "        self.eps = eps\n",
+    "        self.weight = nn.Parameter(torch.ones(dim))\n",
+    "\n",
+    "    def _norm(self, x):\n",
+    "        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        output = self._norm(x.float()).type_as(x)\n",
+    "        return output * self.weight\n",
+    "\n",
+    "def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):\n",
+    "    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n",
+    "    t = torch.arange(end, device=freqs.device, dtype=torch.float32)\n",
+    "    freqs = torch.outer(t, freqs)\n",
+    "    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64\n",
+    "    return freqs_cis\n",
+    "\n",
+    "def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):\n",
+    "    ndim = x.ndim\n",
+    "    assert 0 <= 1 < ndim\n",
+    "    assert freqs_cis.shape == (x.shape[1], x.shape[-1])\n",
+    "    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n",
+    "    return freqs_cis.view(*shape)\n",
+    "\n",
+    "def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):\n",
+    "    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n",
+    "    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n",
+    "    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)\n",
+    "    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)\n",
+    "    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)\n",
+    "    return xq_out.type_as(xq), xk_out.type_as(xk)\n",
+    "\n",
+    "class CausalSelfAttention(nn.Module):\n",
+    "    def __init__(self, config: SmolLMConfig):\n",
+    "        super().__init__()\n",
+    "        self.n_head = config.n_head\n",
+    "        self.n_kv_head = config.n_kv_head\n",
+    "        self.n_embd = config.n_embd\n",
+    "        self.head_dim = config.n_embd // config.n_head\n",
+    "        self.n_rep = self.n_head // self.n_kv_head\n",
+    "\n",
+    "        self.wq = nn.Linear(config.n_embd, config.n_head * self.head_dim, bias=config.bias)\n",
+    "        self.wk = nn.Linear(config.n_embd, config.n_kv_head * self.head_dim, bias=config.bias)\n",
+    "        self.wv = nn.Linear(config.n_embd, config.n_kv_head * self.head_dim, bias=config.bias)\n",
+    "        self.wo = nn.Linear(config.n_head * self.head_dim, config.n_embd, bias=config.bias)\n",
+    "\n",
+    "        self.dropout = config.dropout\n",
+    "        self.resid_dropout = nn.Dropout(config.dropout)\n",
+    "\n",
+    "    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):\n",
+    "        B, T, C = x.shape\n",
+    "\n",
+    "        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)\n",
+    "        xq = xq.view(B, T, self.n_head, self.head_dim)\n",
+    "        xk = xk.view(B, T, self.n_kv_head, self.head_dim)\n",
+    "        xv = xv.view(B, T, self.n_kv_head, self.head_dim)\n",
+    "\n",
+    "        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)\n",
+    "\n",
+    "        # Grouped Query Attention: repeat k/v heads to match q heads\n",
+    "        xk = torch.repeat_interleave(xk, dim=2, repeats=self.n_rep)\n",
+    "        xv = torch.repeat_interleave(xv, dim=2, repeats=self.n_rep)\n",
+    "\n",
+    "        # Make heads batch dimension\n",
+    "        xq = xq.transpose(1, 2)  # (B, n_head, T, head_dim)\n",
+    "        xk = xk.transpose(1, 2)  # (B, n_head, T, head_dim)\n",
+    "        xv = xv.transpose(1, 2)  # (B, n_head, T, head_dim)\n",
+    "\n",
+    "        # Flash Attention\n",
+    "        output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)\n",
+    "\n",
+    "        output = output.transpose(1, 2).contiguous().view(B, T, C)\n",
+    "        return self.resid_dropout(self.wo(output))\n",
+    "\n",
+    "class SwiGLU(nn.Module):\n",
+    "    def __init__(self, config: SmolLMConfig):\n",
+    "        super().__init__()\n",
+    "        self.w1 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias) # Gate\n",
+    "        self.w3 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias) # Value\n",
+    "        self.w2 = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias) # Output\n",
+    "        self.dropout = nn.Dropout(config.dropout)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))\n",
+    "\n",
+    "class Block(nn.Module):\n",
+    "    def __init__(self, config: SmolLMConfig):\n",
+    "        super().__init__()\n",
+    "        self.attention_norm = RMSNorm(config.n_embd, eps=config.rms_norm_eps)\n",
+    "        self.attention = CausalSelfAttention(config)\n",
+    "        self.ffn_norm = RMSNorm(config.n_embd, eps=config.rms_norm_eps)\n",
+    "        self.feed_forward = SwiGLU(config)\n",
+    "\n",
+    "    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor):\n",
+    "        h = x + self.attention(self.attention_norm(x), freqs_cis)\n",
+    "        out = h + self.feed_forward(self.ffn_norm(h))\n",
+    "        return out\n",
+    "\n",
+    "class SmolLM(nn.Module):\n",
+    "    def __init__(self, config: SmolLMConfig):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "        self.tok_embeddings = nn.Embedding(config.vocab_size, config.n_embd)\n",
+    "        self.layers = nn.ModuleList([Block(config) for _ in range(config.n_layer)])\n",
+    "        self.norm = RMSNorm(config.n_embd, eps=config.rms_norm_eps)\n",
+    "        self.output = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n",
+    "\n",
+    "        # Weight sharing\n",
+    "        self.tok_embeddings.weight = self.output.weight\n",
+    "\n",
+    "        # Precompute RoPE frequencies\n",
+    "        self.freqs_cis = precompute_freqs_cis(config.n_embd // config.n_head, config.block_size * 2, config.rope_theta)\n",
+    "\n",
+    "        self.apply(self._init_weights)\n",
+    "\n",
+    "    def _init_weights(self, module):\n",
+    "        if isinstance(module, nn.Linear):\n",
+    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
+    "            if module.bias is not None:\n",
+    "                torch.nn.init.zeros_(module.bias)\n",
+    "        elif isinstance(module, nn.Embedding):\n",
+    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
+    "\n",
+    "    def forward(self, idx, targets=None):\n",
+    "        B, T = idx.shape\n",
+    "        x = self.tok_embeddings(idx)\n",
+    "        \n",
+    "        # Ensure freqs_cis is on the correct device\n",
+    "        if self.freqs_cis.device != x.device:\n",
+    "            self.freqs_cis = self.freqs_cis.to(x.device)\n",
+    "        freqs_cis = self.freqs_cis[:T]\n",
+    "\n",
+    "        for layer in self.layers:\n",
+    "            x = layer(x, freqs_cis)\n",
+    "        \n",
+    "        x = self.norm(x)\n",
+    "        logits = self.output(x)\n",
+    "\n",
+    "        loss = None\n",
+    "        if targets is not None:\n",
+    "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n",
+    "        \n",
+    "        return logits, loss\n",
+    "\n",
+    "# Device selection\n",
+    "device = 'cpu'\n",
+    "if torch.cuda.is_available():\n",
+    "    device = 'cuda'\n",
+    "elif hasattr(torch.backends, \"mps\") and torch.backends.mps.is_available():\n",
+    "    device = \"mps\"\n",
+    "print(f\"using device: {device}\")\n",
+    "\n",
+    "# Data Loading\n",
+    "import tiktoken\n",
+    "\n",
+    "class DataLoaderLite:\n",
+    "    def __init__(self, B, T):\n",
+    "        self.B = B\n",
+    "        self.T = T\n",
+    "\n",
+    "        # Load tokens from disk\n",
+    "        try:\n",
+    "            with open('input.txt', 'r', encoding='utf-8') as f:\n",
+    "                text = f.read()\n",
+    "        except FileNotFoundError:\n",
+    "            print(\"Error: input.txt not found. Please ensure the file exists.\")\n",
+    "            text = \"Hello world \" * 1000 # Fallback for testing if file missing\n",
+    "            \n",
+    "        enc = tiktoken.get_encoding('gpt2') \n",
+    "        tokens = enc.encode(text)\n",
+    "        self.tokens = torch.tensor(tokens)\n",
+    "        print(f'loaded {len(self.tokens)} tokens')\n",
+    "        print(f'1 epoch = {len(self.tokens) // (B * T)} batches')\n",
+    "\n",
+    "        self.current_position = 0\n",
+    "    \n",
+    "    def next_batch(self):\n",
+    "        B, T = self.B, self.T\n",
+    "        buf = self.tokens[self.current_position: self.current_position + B * T + 1]\n",
+    "        x = (buf[:-1]).view(B, T) # inputs\n",
+    "        y = (buf[1:]).view(B, T) # targets\n",
+    "        self.current_position += B*T\n",
+    "        if self.current_position + (B * T + 1) > len(self.tokens):\n",
+    "            self.current_position = 0\n",
+    "        return x, y\n",
+    "\n",
+    "# Training Setup\n",
+    "torch.manual_seed(1337)\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.cuda.manual_seed(1337)\n",
+    "\n",
+    "torch.set_float32_matmul_precision('high')\n",
+    "\n",
+    "config = SmolLMConfig()\n",
+    "model = SmolLM(config)\n",
+    "model.to(device)\n",
+    "\n",
+    "print(f\"Model parameters: {sum(p.numel() for p in model.parameters())/1e6:.2f}M\")\n",
+    "\n",
+    "# Generation Function\n",
+    "@torch.no_grad()\n",
+    "def generate(model, idx, max_new_tokens, temperature=1.0, top_k=None):\n",
+    "    \"\"\"\n",
+    "    Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete\n",
+    "    the sequence max_new_tokens times, feeding the predictions back into the model each time.\n",
+    "    \"\"\"\n",
+    "    for _ in range(max_new_tokens):\n",
+    "        # if the sequence context is growing too long we must crop it at block_size\n",
+    "        idx_cond = idx if idx.size(1) <= model.config.block_size else idx[:, -model.config.block_size:]\n",
+    "        # forward the model to get the logits for the index in the sequence\n",
+    "        logits, _ = model(idx_cond)\n",
+    "        # pluck the logits at the final step and scale by desired temperature\n",
+    "        logits = logits[:, -1, :] / temperature\n",
+    "        # optionally crop the logits to only the top k options\n",
+    "        if top_k is not None:\n",
+    "            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))\n",
+    "            logits[logits < v[:, [-1]]] = -float('Inf')\n",
+    "        # apply softmax to convert logits to (normalized) probabilities\n",
+    "        probs = F.softmax(logits, dim=-1)\n",
+    "        # sample from the distribution\n",
+    "        idx_next = torch.multinomial(probs, num_samples=1)\n",
+    "        # append sampled index to the running sequence and continue\n",
+    "        idx = torch.cat((idx, idx_next), dim=1)\n",
+    "\n",
+    "    return idx\n",
+    "\n",
+    "# Training Loop\n",
+    "train_loader = DataLoaderLite(B = 4, T = 512) # Reduced batch size and context for 4GB GPU\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr = 3e-4)\n",
+    "\n",
+    "import time\n",
+    "import os\n",
+    "\n",
+    "max_steps = 5000\n",
+    "eval_interval = 500\n",
+    "save_path = \"smollm_135_checkpoint.pth\"\n",
+    "\n",
+    "print(\"Starting training...\")\n",
+    "for i in range(max_steps):\n",
+    "    t0 = time.time()\n",
+    "    x, y = train_loader.next_batch()\n",
+    "    x, y = x.to(device), y.to(device)\n",
+    "    optimizer.zero_grad()\n",
+    "    \n",
+    "    # Mixed precision training\n",
+    "    with torch.autocast(device_type=device, dtype=torch.bfloat16 if device=='cuda' else torch.float32):\n",
+    "        logits, loss = model(x, y) \n",
+    "    \n",
+    "    loss.backward()\n",
+    "    optimizer.step()\n",
+    "    \n",
+    "    if device == 'cuda':\n",
+    "        torch.cuda.synchronize() \n",
+    "        \n",
+    "    t1 = time.time()\n",
+    "    dt = (t1 - t0) * 1000\n",
+    "    tokens_per_sec = (train_loader.B * train_loader.T) / (t1 - t0)\n",
+    "    \n",
+    "    if i % 10 == 0:\n",
+    "        print(f'step {i} | loss: {loss.item():.4f} | dt: {dt:.2f}ms | tok/sec: {tokens_per_sec:.2f}')\n",
+    "        \n",
+    "    # Generate output every 500 steps\n",
+    "    if i > 0 and i % eval_interval == 0:\n",
+    "        print(f\"\\n--- Generating text at step {i} ---\")\n",
+    "        context = torch.zeros((1, 1), dtype=torch.long, device=device) # Start with token 0 (usually valid)\n",
+    "        generated = generate(model, context, max_new_tokens=50)\n",
+    "        # Decode using tiktoken (gpt2 encoding as used in DataLoader)\n",
+    "        enc = tiktoken.get_encoding('gpt2')\n",
+    "        decoded = enc.decode(generated[0].tolist())\n",
+    "        # Force ASCII for Windows console compatibility\n",
+    "        print(decoded.encode('ascii', errors='ignore').decode('ascii'))\n",
+    "        print(\"-----------------------------------\\n\")\n",
+    "\n",
+    "# Save checkpoint\n",
+    "print(f\"Saving model to {save_path}\")\n",
+    "torch.save(model.state_dict(), save_path)\n",
+    "\n",
+    "# Resume training demonstration\n",
+    "print(\"\\n--- Resuming training from checkpoint ---\")\n",
+    "# Re-initialize model to prove loading works\n",
+    "model_new = SmolLM(config)\n",
+    "model_new.to(device)\n",
+    "model_new.load_state_dict(torch.load(save_path))\n",
+    "print(\"Checkpoint loaded successfully.\")\n",
+    "\n",
+    "optimizer_new = torch.optim.AdamW(model_new.parameters(), lr = 3e-4)\n",
+    "\n",
+    "# Train for another 50 steps\n",
+    "for i in range(50):\n",
+    "    t0 = time.time()\n",
+    "    x, y = train_loader.next_batch()\n",
+    "    x, y = x.to(device), y.to(device)\n",
+    "    optimizer_new.zero_grad()\n",
+    "    \n",
+    "    with torch.autocast(device_type=device, dtype=torch.bfloat16 if device=='cuda' else torch.float32):\n",
+    "        logits, loss = model_new(x, y) \n",
+    "    \n",
+    "    loss.backward()\n",
+    "    optimizer_new.step()\n",
+    "    \n",
+    "    if device == 'cuda':\n",
+    "        torch.cuda.synchronize()\n",
+    "        \n",
+    "    if i % 10 == 0:\n",
+    "        print(f'Resume step {i} | loss: {loss.item():.4f}')\n",
+    "\n",
+    "print(\"Resumed training completed.\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a720457f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.13.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

app.py

@@ -0,0 +1,160 @@
+"""
+Gradio App for Sentence Completion
+Main entry point for Hugging Face Spaces
+"""
+
+import gradio as gr
+import torch
+from inference import load_model, generate_text, get_device
+
+
+# Global model variable
+model = None
+device = None
+
+
+def initialize_model(model_path=None):
+    """Initialize the model on startup"""
+    global model, device
+    try:
+        model, device = load_model(model_path=model_path)
+        #model.eval()  # Set to eval mode once
+        return f"Model loaded successfully on device: {device}"
+    except Exception as e:
+        return f"Error loading model: {str(e)}"
+
+
+def complete_sentence(prompt, max_tokens, top_k, temperature):
+    """Generate sentence completion based on prompt"""
+    global model, device
+    
+    if model is None:
+        return "Error: Model not loaded. Please restart the app."
+    
+    if not prompt.strip():
+        return "Please enter a prompt to complete."
+    
+    try:
+        # Generate completion
+        print(prompt)
+        generated_text = generate_text(
+            prompt=prompt,
+            model=model,
+            max_tokens=int(max_tokens),
+            device=device
+        )
+        print(device)
+        
+        return generated_text
+    except Exception as e:
+        return f"Error generating text: {str(e)}"
+
+
+def create_interface():
+    """Create and return the Gradio interface"""
+    
+    # Initialize model on startup
+    # Try to load from common checkpoint paths
+    checkpoint_paths = [
+        './model/smollm_135_checkpoint.pth',
+        './model/model.pth',
+        'model.pt',
+        'checkpoint.pth',
+    ]
+    
+    model_path = None
+    for path in checkpoint_paths:
+        import os
+        if os.path.exists(path):
+            model_path = path
+            print(f"Model found at {path}")
+            break
+        else:
+            print(f"Model not found at {path}")
+    
+    status = initialize_model(model_path=model_path)
+    print(status)
+    
+    # Create Gradio interface
+    with gr.Blocks(title="Sentence Completion with SmolLM-135M") as demo:
+        gr.Markdown(
+            """
+            # Sentence Completion with SmolLM-135M
+            
+            Enter a prompt and the model will complete the sentence for you.
+            Adjust the parameters to control the generation behavior.
+            """
+        )
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter your prompt here...",
+                    lines=3,
+                    value="The future of artificial intelligence is"
+                )
+                
+                with gr.Row():
+                    max_tokens_slider = gr.Slider(
+                        minimum=10,
+                        maximum=200,
+                        value=50,
+                        step=10,
+                        label="Max Tokens"
+                    )
+                    
+                    top_k_slider = gr.Slider(
+                        minimum=1,
+                        maximum=100,
+                        value=50,
+                        step=1,
+                        label="Top-K"
+                    )
+                    
+                    temperature_slider = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=1.0,
+                        step=0.1,
+                        label="Temperature"
+                    )
+                
+                generate_btn = gr.Button("Generate", variant="primary")
+            
+            with gr.Column(scale=2):
+                output_text = gr.Textbox(
+                    label="Generated Text",
+                    lines=10,
+                    interactive=False
+                )
+        
+        gr.Markdown(
+            """
+            ### Parameters:
+            - **Max Tokens**: Maximum number of tokens to generate
+            - **Top-K**: Sample from top K most likely tokens (lower = more focused)
+            - **Temperature**: Controls randomness (lower = more deterministic, higher = more creative)
+            """
+        )
+        
+        # Set up the generate function
+        generate_btn.click(
+            fn=complete_sentence,
+            inputs=[prompt_input, max_tokens_slider, top_k_slider, temperature_slider],
+            outputs=output_text
+        )
+        
+        # Also generate on Enter key press
+        prompt_input.submit(
+            fn=complete_sentence,
+            inputs=[prompt_input, max_tokens_slider, top_k_slider, temperature_slider],
+            outputs=output_text
+        )
+    
+    return demo
+
+
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch(share=False)

check_cuda.py

@@ -0,0 +1,9 @@
+import torch
+
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"CUDA device count: {torch.cuda.device_count()}")
+if torch.cuda.is_available():
+    print(f"Current device: {torch.cuda.current_device()}")
+    print(f"Device name: {torch.cuda.get_device_name(0)}")
+else:
+    print("CUDA not available")

checkpoint_info.txt

@@ -0,0 +1,2 @@
+Checkpoint type: <class 'collections.OrderedDict'>
+Keys: ['tok_embeddings.weight', 'layers.0.attention_norm.weight', 'layers.0.attention.wq.weight', 'layers.0.attention.wk.weight', 'layers.0.attention.wv.weight', 'layers.0.attention.wo.weight', 'layers.0.ffn_norm.weight', 'layers.0.feed_forward.w1.weight', 'layers.0.feed_forward.w3.weight', 'layers.0.feed_forward.w2.weight', 'layers.1.attention_norm.weight', 'layers.1.attention.wq.weight', 'layers.1.attention.wk.weight', 'layers.1.attention.wv.weight', 'layers.1.attention.wo.weight', 'layers.1.ffn_norm.weight', 'layers.1.feed_forward.w1.weight', 'layers.1.feed_forward.w3.weight', 'layers.1.feed_forward.w2.weight', 'layers.2.attention_norm.weight', 'layers.2.attention.wq.weight', 'layers.2.attention.wk.weight', 'layers.2.attention.wv.weight', 'layers.2.attention.wo.weight', 'layers.2.ffn_norm.weight', 'layers.2.feed_forward.w1.weight', 'layers.2.feed_forward.w3.weight', 'layers.2.feed_forward.w2.weight', 'layers.3.attention_norm.weight', 'layers.3.attention.wq.weight', 'layers.3.attention.wk.weight', 'layers.3.attention.wv.weight', 'layers.3.attention.wo.weight', 'layers.3.ffn_norm.weight', 'layers.3.feed_forward.w1.weight', 'layers.3.feed_forward.w3.weight', 'layers.3.feed_forward.w2.weight', 'layers.4.attention_norm.weight', 'layers.4.attention.wq.weight', 'layers.4.attention.wk.weight', 'layers.4.attention.wv.weight', 'layers.4.attention.wo.weight', 'layers.4.ffn_norm.weight', 'layers.4.feed_forward.w1.weight', 'layers.4.feed_forward.w3.weight', 'layers.4.feed_forward.w2.weight', 'layers.5.attention_norm.weight', 'layers.5.attention.wq.weight', 'layers.5.attention.wk.weight', 'layers.5.attention.wv.weight', 'layers.5.attention.wo.weight', 'layers.5.ffn_norm.weight', 'layers.5.feed_forward.w1.weight', 'layers.5.feed_forward.w3.weight', 'layers.5.feed_forward.w2.weight', 'layers.6.attention_norm.weight', 'layers.6.attention.wq.weight', 'layers.6.attention.wk.weight', 'layers.6.attention.wv.weight', 'layers.6.attention.wo.weight', 'layers.6.ffn_norm.weight', 'layers.6.feed_forward.w1.weight', 'layers.6.feed_forward.w3.weight', 'layers.6.feed_forward.w2.weight', 'layers.7.attention_norm.weight', 'layers.7.attention.wq.weight', 'layers.7.attention.wk.weight', 'layers.7.attention.wv.weight', 'layers.7.attention.wo.weight', 'layers.7.ffn_norm.weight', 'layers.7.feed_forward.w1.weight', 'layers.7.feed_forward.w3.weight', 'layers.7.feed_forward.w2.weight', 'layers.8.attention_norm.weight', 'layers.8.attention.wq.weight', 'layers.8.attention.wk.weight', 'layers.8.attention.wv.weight', 'layers.8.attention.wo.weight', 'layers.8.ffn_norm.weight', 'layers.8.feed_forward.w1.weight', 'layers.8.feed_forward.w3.weight', 'layers.8.feed_forward.w2.weight', 'layers.9.attention_norm.weight', 'layers.9.attention.wq.weight', 'layers.9.attention.wk.weight', 'layers.9.attention.wv.weight', 'layers.9.attention.wo.weight', 'layers.9.ffn_norm.weight', 'layers.9.feed_forward.w1.weight', 'layers.9.feed_forward.w3.weight', 'layers.9.feed_forward.w2.weight', 'layers.10.attention_norm.weight', 'layers.10.attention.wq.weight', 'layers.10.attention.wk.weight', 'layers.10.attention.wv.weight', 'layers.10.attention.wo.weight', 'layers.10.ffn_norm.weight', 'layers.10.feed_forward.w1.weight', 'layers.10.feed_forward.w3.weight', 'layers.10.feed_forward.w2.weight', 'layers.11.attention_norm.weight', 'layers.11.attention.wq.weight', 'layers.11.attention.wk.weight', 'layers.11.attention.wv.weight', 'layers.11.attention.wo.weight', 'layers.11.ffn_norm.weight', 'layers.11.feed_forward.w1.weight', 'layers.11.feed_forward.w3.weight', 'layers.11.feed_forward.w2.weight', 'layers.12.attention_norm.weight', 'layers.12.attention.wq.weight', 'layers.12.attention.wk.weight', 'layers.12.attention.wv.weight', 'layers.12.attention.wo.weight', 'layers.12.ffn_norm.weight', 'layers.12.feed_forward.w1.weight', 'layers.12.feed_forward.w3.weight', 'layers.12.feed_forward.w2.weight', 'layers.13.attention_norm.weight', 'layers.13.attention.wq.weight', 'layers.13.attention.wk.weight', 'layers.13.attention.wv.weight', 'layers.13.attention.wo.weight', 'layers.13.ffn_norm.weight', 'layers.13.feed_forward.w1.weight', 'layers.13.feed_forward.w3.weight', 'layers.13.feed_forward.w2.weight', 'layers.14.attention_norm.weight', 'layers.14.attention.wq.weight', 'layers.14.attention.wk.weight', 'layers.14.attention.wv.weight', 'layers.14.attention.wo.weight', 'layers.14.ffn_norm.weight', 'layers.14.feed_forward.w1.weight', 'layers.14.feed_forward.w3.weight', 'layers.14.feed_forward.w2.weight', 'layers.15.attention_norm.weight', 'layers.15.attention.wq.weight', 'layers.15.attention.wk.weight', 'layers.15.attention.wv.weight', 'layers.15.attention.wo.weight', 'layers.15.ffn_norm.weight', 'layers.15.feed_forward.w1.weight', 'layers.15.feed_forward.w3.weight', 'layers.15.feed_forward.w2.weight', 'layers.16.attention_norm.weight', 'layers.16.attention.wq.weight', 'layers.16.attention.wk.weight', 'layers.16.attention.wv.weight', 'layers.16.attention.wo.weight', 'layers.16.ffn_norm.weight', 'layers.16.feed_forward.w1.weight', 'layers.16.feed_forward.w3.weight', 'layers.16.feed_forward.w2.weight', 'layers.17.attention_norm.weight', 'layers.17.attention.wq.weight', 'layers.17.attention.wk.weight', 'layers.17.attention.wv.weight', 'layers.17.attention.wo.weight', 'layers.17.ffn_norm.weight', 'layers.17.feed_forward.w1.weight', 'layers.17.feed_forward.w3.weight', 'layers.17.feed_forward.w2.weight', 'layers.18.attention_norm.weight', 'layers.18.attention.wq.weight', 'layers.18.attention.wk.weight', 'layers.18.attention.wv.weight', 'layers.18.attention.wo.weight', 'layers.18.ffn_norm.weight', 'layers.18.feed_forward.w1.weight', 'layers.18.feed_forward.w3.weight', 'layers.18.feed_forward.w2.weight', 'layers.19.attention_norm.weight', 'layers.19.attention.wq.weight', 'layers.19.attention.wk.weight', 'layers.19.attention.wv.weight', 'layers.19.attention.wo.weight', 'layers.19.ffn_norm.weight', 'layers.19.feed_forward.w1.weight', 'layers.19.feed_forward.w3.weight', 'layers.19.feed_forward.w2.weight', 'layers.20.attention_norm.weight', 'layers.20.attention.wq.weight', 'layers.20.attention.wk.weight', 'layers.20.attention.wv.weight', 'layers.20.attention.wo.weight', 'layers.20.ffn_norm.weight', 'layers.20.feed_forward.w1.weight', 'layers.20.feed_forward.w3.weight', 'layers.20.feed_forward.w2.weight', 'layers.21.attention_norm.weight', 'layers.21.attention.wq.weight', 'layers.21.attention.wk.weight', 'layers.21.attention.wv.weight', 'layers.21.attention.wo.weight', 'layers.21.ffn_norm.weight', 'layers.21.feed_forward.w1.weight', 'layers.21.feed_forward.w3.weight', 'layers.21.feed_forward.w2.weight', 'layers.22.attention_norm.weight', 'layers.22.attention.wq.weight', 'layers.22.attention.wk.weight', 'layers.22.attention.wv.weight', 'layers.22.attention.wo.weight', 'layers.22.ffn_norm.weight', 'layers.22.feed_forward.w1.weight', 'layers.22.feed_forward.w3.weight', 'layers.22.feed_forward.w2.weight', 'layers.23.attention_norm.weight', 'layers.23.attention.wq.weight', 'layers.23.attention.wk.weight', 'layers.23.attention.wv.weight', 'layers.23.attention.wo.weight', 'layers.23.ffn_norm.weight', 'layers.23.feed_forward.w1.weight', 'layers.23.feed_forward.w3.weight', 'layers.23.feed_forward.w2.weight', 'layers.24.attention_norm.weight', 'layers.24.attention.wq.weight', 'layers.24.attention.wk.weight', 'layers.24.attention.wv.weight', 'layers.24.attention.wo.weight', 'layers.24.ffn_norm.weight', 'layers.24.feed_forward.w1.weight', 'layers.24.feed_forward.w3.weight', 'layers.24.feed_forward.w2.weight', 'layers.25.attention_norm.weight', 'layers.25.attention.wq.weight', 'layers.25.attention.wk.weight', 'layers.25.attention.wv.weight', 'layers.25.attention.wo.weight', 'layers.25.ffn_norm.weight', 'layers.25.feed_forward.w1.weight', 'layers.25.feed_forward.w3.weight', 'layers.25.feed_forward.w2.weight', 'layers.26.attention_norm.weight', 'layers.26.attention.wq.weight', 'layers.26.attention.wk.weight', 'layers.26.attention.wv.weight', 'layers.26.attention.wo.weight', 'layers.26.ffn_norm.weight', 'layers.26.feed_forward.w1.weight', 'layers.26.feed_forward.w3.weight', 'layers.26.feed_forward.w2.weight', 'layers.27.attention_norm.weight', 'layers.27.attention.wq.weight', 'layers.27.attention.wk.weight', 'layers.27.attention.wv.weight', 'layers.27.attention.wo.weight', 'layers.27.ffn_norm.weight', 'layers.27.feed_forward.w1.weight', 'layers.27.feed_forward.w3.weight', 'layers.27.feed_forward.w2.weight', 'layers.28.attention_norm.weight', 'layers.28.attention.wq.weight', 'layers.28.attention.wk.weight', 'layers.28.attention.wv.weight', 'layers.28.attention.wo.weight', 'layers.28.ffn_norm.weight', 'layers.28.feed_forward.w1.weight', 'layers.28.feed_forward.w3.weight', 'layers.28.feed_forward.w2.weight', 'layers.29.attention_norm.weight', 'layers.29.attention.wq.weight', 'layers.29.attention.wk.weight', 'layers.29.attention.wv.weight', 'layers.29.attention.wo.weight', 'layers.29.ffn_norm.weight', 'layers.29.feed_forward.w1.weight', 'layers.29.feed_forward.w3.weight', 'layers.29.feed_forward.w2.weight', 'norm.weight', 'output.weight']

inference.py

@@ -0,0 +1,126 @@
+"""
+Inference and Model Loading Utilities
+"""
+
+import os
+import torch
+from torch.nn import functional as F
+import tiktoken
+from model import SmolLM, SmolLMConfig
+
+
+def get_device():
+    """Auto-detect and return the best available device"""
+    print(f"[DEBUG] Checking device availability...")
+    print(f"[DEBUG] torch.cuda.is_available(): {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"[DEBUG] CUDA device: {torch.cuda.get_device_name(0)}")
+        return 'cuda'
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        print(f"[DEBUG] Using MPS device")
+        return "mps"
+    else:
+        print(f"[DEBUG] Falling back to CPU")
+        return 'cpu'
+
+
+def load_model(model_path=None, device=None):
+    """
+    Load SmolLM model from checkpoint.
+    
+    Args:
+        model_path: Path to saved model checkpoint (.pth or .pt file)
+        device: Device to load model on (auto-detected if None)
+    
+    Returns:
+        Loaded model and device
+    """
+    if device is None:
+        device = get_device()
+    
+    # Try to load saved checkpoint first
+    if model_path and os.path.exists(model_path):
+        try:
+            print(f"Loading saved model from {model_path}...")
+            model = SmolLM.load_checkpoint(model_path, device=device)
+            return model, device
+        except Exception as e:
+            print(f"Failed to load saved model: {e}")
+    
+    # Fallback to untrained model
+    print("Creating model with default config (untrained)...")
+    config = SmolLMConfig()
+    model = SmolLM(config)
+    model.to(device)
+    return model, device
+
+
+def generate_text(prompt, model, max_tokens=50, top_k=50, temperature=1.0, device="cpu"):
+    """
+    Generate text completion for a given prompt using the SmolLM model.
+    
+    Args:
+        prompt: Input text prompt
+        model: SmolLM model instance
+        max_tokens: Maximum number of tokens to generate
+        top_k: Top-k sampling parameter (None for no top-k filtering)
+        temperature: Temperature for sampling (higher = more random)
+        device: Device to run inference on
+    
+    Returns:
+        Generated text string (including original prompt)
+    """
+# Global tokenizer cache
+_TOKENIZER = None
+
+def _get_tokenizer():
+    global _TOKENIZER
+    if _TOKENIZER is None:
+        _TOKENIZER = tiktoken.get_encoding("gpt2")
+    return _TOKENIZER
+
+def generate_text(prompt, model, max_tokens=50, top_k=50, temperature=1.0, device="cpu"):
+    """
+    Generate text completion for a given prompt using the SmolLM model.
+    
+    Args:
+        prompt: Input text prompt
+        model: SmolLM model instance
+        max_tokens: Maximum number of tokens to generate
+        top_k: Top-k sampling parameter (None for no top-k filtering)
+        temperature: Temperature for sampling (higher = more random)
+        device: Device to run inference on
+    
+    Returns:
+        Generated text string (including original prompt)
+    """
+    enc = _get_tokenizer()
+    model.eval()
+    
+    with torch.no_grad():
+        # tokenize prompt
+        input_ids = enc.encode(prompt)
+        x = torch.tensor(input_ids, dtype=torch.long, device=device).unsqueeze(0)
+        print(max_tokens)
+        past_kv = None
+        generated_ids = list(input_ids)
+
+        for _ in range(max_tokens):
+            print(x.shape)
+            logits, _ = model(x)
+            print(logits.shape)
+            logits = logits[:, -1, :] / temperature
+            print(logits.shape) 
+            if top_k is not None:
+                topk = torch.topk(logits, top_k, dim=-1)
+                mask = logits < topk.values[:, -1].unsqueeze(-1)
+                logits = logits.masked_fill(mask, -float("inf"))
+
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            x = torch.cat((x, next_token), dim=1)
+
+        generated_ids = x[0].tolist()
+
+        # generated_ids already contains the prompt and generated tokens
+    return enc.decode(generated_ids)

inspect_checkpoint.py

@@ -0,0 +1,19 @@
+import torch
+import sys
+
+def inspect_checkpoint():
+    path = './model/smollm_135_checkpoint.pth'
+    output_file = 'checkpoint_info.txt'
+    with open(output_file, 'w') as f:
+        try:
+            checkpoint = torch.load(path, map_location='cpu')
+            f.write(f"Checkpoint type: {type(checkpoint)}\n")
+            if isinstance(checkpoint, dict):
+                f.write(f"Keys: {list(checkpoint.keys())}\n")
+            else:
+                f.write("Checkpoint is not a dictionary.\n")
+        except Exception as e:
+            f.write(f"Error loading checkpoint: {e}\n")
+
+if __name__ == "__main__":
+    inspect_checkpoint()

main.py

@@ -0,0 +1,6 @@
+def main():
+    print("Hello from smollm-135!")
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,239 @@
+"""
+SmolLM-135M Implementation (Llama Architecture)
+Based on: https://huggingface.co/HuggingFaceTB/SmolLM-135M
+"""
+
+import math
+import inspect
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+# Configuration for SmolLM-135M
+@dataclass
+class SmolLMConfig:
+    block_size: int = 512 # Reduced to 512 for 4GB GPU training
+    vocab_size: int = 50304 # Aligned to 50304 for tiktoken compatibility (SmolLM native is 49152)
+    n_layer: int = 30
+    n_head: int = 9
+    n_kv_head: int = 3 # Grouped Query Attention (GQA)
+    n_embd: int = 576
+    intermediate_size: int = 1536 # SwiGLU intermediate size
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+    dropout: float = 0.0
+    bias: bool = False # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: SmolLMConfig):
+        super().__init__()
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.n_embd // config.n_head
+        self.n_rep = self.n_head // self.n_kv_head
+
+        self.wq = nn.Linear(config.n_embd, config.n_head * self.head_dim, bias=config.bias)
+        self.wk = nn.Linear(config.n_embd, config.n_kv_head * self.head_dim, bias=config.bias)
+        self.wv = nn.Linear(config.n_embd, config.n_kv_head * self.head_dim, bias=config.bias)
+        self.wo = nn.Linear(config.n_head * self.head_dim, config.n_embd, bias=config.bias)
+
+        self.dropout = config.dropout
+        self.resid_dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, C = x.shape
+
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        xq = xq.view(B, T, self.n_head, self.head_dim)
+        xk = xk.view(B, T, self.n_kv_head, self.head_dim)
+        xv = xv.view(B, T, self.n_kv_head, self.head_dim)
+
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+
+        # Grouped Query Attention: repeat k/v heads to match q heads
+        xk = torch.repeat_interleave(xk, dim=2, repeats=self.n_rep)
+        xv = torch.repeat_interleave(xv, dim=2, repeats=self.n_rep)
+        
+        # Make heads batch dimension
+        xq = xq.transpose(1, 2)  # (B, n_head, T, head_dim)
+        xk = xk.transpose(1, 2)  # (B, n_head, T, head_dim)
+        xv = xv.transpose(1, 2)  # (B, n_head, T, head_dim)
+        
+        if past_kv is not None:
+            k_cache, v_cache = past_kv
+            xk = torch.cat([k_cache, xk], dim=2)
+            xv = torch.cat([v_cache, xv], dim=2)
+            
+        current_kv = (xk, xv)
+
+        # Flash Attention
+        if past_kv is not None and T == 1:
+            # Optimization: no causal mask needed for the last token attending to all previous
+            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=False)
+        else:
+            output = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
+
+        output = output.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_dropout(self.wo(output)), current_kv
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: SmolLMConfig):
+        super().__init__()
+        self.w1 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias) # Gate
+        self.w3 = nn.Linear(config.n_embd, config.intermediate_size, bias=config.bias) # Value
+        self.w2 = nn.Linear(config.intermediate_size, config.n_embd, bias=config.bias) # Output
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+class Block(nn.Module):
+    def __init__(self, config: SmolLMConfig):
+        super().__init__()
+        self.attention_norm = RMSNorm(config.n_embd, eps=config.rms_norm_eps)
+        self.attention = CausalSelfAttention(config)
+        self.ffn_norm = RMSNorm(config.n_embd, eps=config.rms_norm_eps)
+        self.feed_forward = SwiGLU(config)
+
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, layer_kv = self.attention(self.attention_norm(x), freqs_cis, past_kv)
+        h = x + attn_out
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out, layer_kv
+
+class SmolLM(nn.Module):
+    def __init__(self, config: SmolLMConfig):
+        super().__init__()
+        self.config = config
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.n_embd)
+        self.layers = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
+        self.norm = RMSNorm(config.n_embd, eps=config.rms_norm_eps)
+        self.output = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+
+        # Weight sharing
+        self.tok_embeddings.weight = self.output.weight
+
+        # Precompute RoPE frequencies
+        self.freqs_cis = precompute_freqs_cis(config.n_embd // config.n_head, config.block_size * 2, config.rope_theta)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None, past_kv=None):
+        B, T = idx.shape
+        x = self.tok_embeddings(idx)
+        
+        # Determine starting position for RoPE
+        start_pos = 0
+        
+            
+        # Ensure freqs_cis is on the correct device
+        if self.freqs_cis.device != x.device:
+            self.freqs_cis = self.freqs_cis.to(x.device)
+        
+        # Select freqs_cis for the current positions
+        freqs_cis = self.freqs_cis[start_pos : start_pos + T]
+
+        new_past_kv = []
+        for i, layer in enumerate(self.layers):
+            layer_past_kv = past_kv[i] if past_kv is not None else None
+            x, layer_kv = layer(x, freqs_cis, past_kv=layer_past_kv)
+            new_past_kv.append(layer_kv)
+        
+        x = self.norm(x)
+        logits = self.output(x)
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+            
+        return logits, loss
+
+    def save_checkpoint(self, filepath):
+        """Save model checkpoint with config"""
+        checkpoint = {
+            'model_state_dict': self.state_dict(),
+            'config': {
+                'block_size': self.config.block_size,
+                'vocab_size': self.config.vocab_size,
+                'n_layer': self.config.n_layer,
+                'n_head': self.config.n_head,
+                'n_kv_head': self.config.n_kv_head,
+                'n_embd': self.config.n_embd,
+                'intermediate_size': self.config.intermediate_size,
+                'rms_norm_eps': self.config.rms_norm_eps,
+                'rope_theta': self.config.rope_theta,
+                'dropout': self.config.dropout,
+                'bias': self.config.bias
+            }
+        }
+        torch.save(checkpoint, filepath)
+        print(f"Model saved to {filepath}")
+
+    @classmethod
+    def load_checkpoint(cls, filepath, device='cpu'):
+        """Load model from checkpoint file"""
+        checkpoint = torch.load(filepath, map_location=device)
+        
+        if isinstance(checkpoint, dict) and 'config' in checkpoint:
+            # Checkpoint contains config and state dict
+            config_dict = checkpoint['config']
+            config = SmolLMConfig(**config_dict)
+            state_dict = checkpoint['model_state_dict']
+        else:
+            # Checkpoint is likely just the state dict
+            print("Warning: Checkpoint does not contain config. Using default SmolLMConfig.")
+            config = SmolLMConfig()
+            state_dict = checkpoint
+
+        model = cls(config)
+        model.load_state_dict(state_dict)
+        model.to(device)
+        print(f"Model loaded from {filepath}")
+        return model

model/smollm_135_checkpoint.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:691cbc7dafd54bde47470447eb33f27ab2f6ac5e77cd45b71942aceeac678b85
+size 540826259

profile_app.py

@@ -0,0 +1,33 @@
+import time
+import torch
+from inference import load_model, generate_text, get_device
+
+def profile_app():
+    print("Profiling app performance...")
+    
+    # 1. Measure Model Loading
+    start_load = time.time()
+    checkpoint_path = './model/smollm_135_checkpoint.pth'
+    model, device = load_model(model_path=checkpoint_path)
+    print(device)
+    end_load = time.time()
+    print(f"Model loading time: {end_load - start_load:.4f}s")
+    print(f"Device: {device}")
+    
+    # 2. Measure Generation
+    prompt = "The future of AI is"
+    max_tokens = 50
+    
+    start_gen = time.time()
+    output = generate_text(prompt, model, max_tokens=max_tokens, device=device)
+    end_gen = time.time()
+    
+    duration = end_gen - start_gen
+    tokens_per_sec = max_tokens / duration
+    
+    print(f"Generation time: {duration:.4f}s")
+    print(f"Tokens per second: {tokens_per_sec:.2f}")
+    print(f"Output: {output}")
+
+if __name__ == "__main__":
+    profile_app()

pyproject.toml

@@ -0,0 +1,12 @@
+[project]
+name = "smollm-135"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "gradio>=6.0.1",
+    "tiktoken>=0.12.0",
+    "torch>=2.9.1",
+    "transformers>=4.57.3",
+]

requirements.txt

@@ -0,0 +1,4 @@
+gradio>=4.0.0
+torch==2.0.0
+transformers>=4.30.0
+tiktoken>=0.5.0

test_inference.py

@@ -0,0 +1,25 @@
+from inference import load_model, generate_text
+import os
+
+def test_inference():
+    print("Testing inference...")
+    
+    checkpoint_path = './model/smollm_135_checkpoint.pth'
+    if not os.path.exists(checkpoint_path):
+        print(f"Warning: Checkpoint not found at {checkpoint_path}")
+        checkpoint_path = None
+    
+    # Load model
+    model, device = load_model(model_path=checkpoint_path)
+    print(f"Model loaded on {device}")
+    
+    # Generate text
+    prompt = "Hello, world"
+    print(f"Generating text for prompt: '{prompt}'")
+    generated = generate_text(prompt, model, max_tokens=10, device=device)
+    print(f"Generated text: {generated}")
+    
+    print("Inference test passed.")
+
+if __name__ == "__main__":
+    test_inference()

test_kv_cache.py

@@ -0,0 +1,45 @@
+import torch
+import time
+from model import SmolLM, SmolLMConfig
+from inference import load_model
+
+def test_kv_performance():
+    print("Testing KV cache performance...")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"Device: {device}")
+    
+    # Load model (untrained is fine for performance test)
+    config = SmolLMConfig()
+    model = SmolLM(config).to(device)
+    model.eval()
+    
+    prompt_len = 10
+    gen_len = 50
+    input_ids = torch.randint(0, config.vocab_size, (1, prompt_len)).to(device)
+    
+    # Measure generation with KV cache
+    start_time = time.time()
+    past_kv = None
+    x = input_ids
+    
+    with torch.no_grad():
+        # Prefill
+        _, _, past_kv = model(x)
+        
+        # Generate
+        for _ in range(gen_len):
+            model_input = x[:, -1:]
+            _, _, past_kv = model(model_input, past_kv=past_kv)
+            # Dummy token selection
+            next_token = torch.tensor([[0]], device=device)
+            x = torch.cat([x, next_token], dim=1)
+            
+    end_time = time.time()
+    duration = end_time - start_time
+    tokens_per_sec = gen_len / duration
+    
+    print(f"Generated {gen_len} tokens in {duration:.4f}s")
+    print(f"Speed: {tokens_per_sec:.2f} tokens/sec")
+
+if __name__ == "__main__":
+    test_kv_performance()

test_model.py

@@ -0,0 +1,22 @@
+import torch
+from model import SmolLM, SmolLMConfig
+
+def test_model():
+    print("Testing SmolLM model...")
+    config = SmolLMConfig()
+    model = SmolLM(config)
+    print("Model instantiated successfully.")
+    
+    # Create dummy input
+    idx = torch.randint(0, config.vocab_size, (1, 128))
+    print(f"Input shape: {idx.shape}")
+    
+    # Forward pass
+    logits, loss, _ = model(idx)
+    print(f"Logits shape: {logits.shape}")
+    
+    assert logits.shape == (1, 128, config.vocab_size)
+    print("Forward pass successful.")
+
+if __name__ == "__main__":
+    test_model()

test_tiktoken.py

@@ -0,0 +1,16 @@
+import time
+import tiktoken
+
+def test_tiktoken():
+    start = time.time()
+    enc = tiktoken.get_encoding("gpt2")
+    end = time.time()
+    print(f"tiktoken load time: {end - start:.4f}s")
+    
+    start = time.time()
+    enc.encode("Hello world")
+    end = time.time()
+    print(f"encoding time: {end - start:.4f}s")
+
+if __name__ == "__main__":
+    test_tiktoken()