We will be using the [Shakespeare dataset](https://github.com/saforem2/wordplay/blob/main/data/shakespeare/readme.md) to train a (~ small) 10M param LLM _from scratch_.

<div>

<div align="center" style="text-align:center;">

<img src="https://github.com/saforem2/wordplay/blob/main/assets/shakespeare.jpeg?raw=true" width="45%" align="center" /><br>

Image generated from [stabilityai/stable-diffusion](https://huggingface.co/spaces/stabilityai/stable-diffusion) on [🤗 Spaces](https://huggingface.co/spaces).<br>

</div>

<details closed><summary>Prompt Details</summary>

<ul>
<li>Prompt:</li>
<t><q>
Shakespeare himself, dressed in full Shakespearean garb,
writing code at a modern workstation with multiple monitors, hacking away profusely,
backlit, high quality for publication
</q></t>

<li>Negative Prompt:</li>
<t><q>
low quality, 3d, photorealistic, ugly
</q></t>
</ul>

</details>

</div>


## Install / Setup

<div class="alert alert-block alert-warning">
<b>Warning!</b><br>  

**IF YOU ARE EXECUTING ON GOOGLE COLAB**:  

You will need to restart your runtime (`Runtime` $\rightarrow\,$ `Restart runtime`)  
_after_ executing the following cell:

</div>

```{python}
#| colab: {base_uri: https://localhost:8080/}
#| execution: {iopub.execute_input: '2023-11-30T16:35:09.547786Z', iopub.status.busy: '2023-11-30T16:35:09.547243Z', iopub.status.idle: '2023-11-30T16:35:09.697821Z', shell.execute_reply: '2023-11-30T16:35:09.697442Z', shell.execute_reply.started: '2023-11-30T16:35:09.547769Z'}
#| jupyter: {outputs_hidden: false, source_hidden: false}
%%bash

python3 -c 'import wordplay; print(wordplay.__file__)' 2> '/dev/null'

if [[ $? -eq 0 ]]; then
    echo "Has wordplay installed. Nothing to do."
else
    echo "Does not have wordplay installed. Installing..."
    git clone 'https://github.com/saforem2/wordplay'
    python3 wordplay/data/shakespeare_char/prepare.py
    python3 wordplay/data/shakespeare/prepare.py
    python3 -m pip install deepspeed
    python3 -m pip install -e wordplay
fi
```

## Post Install

If installed correctly, you should be able to:

```python
>>> import wordplay
>>> wordplay.__file__
'/path/to/wordplay/src/wordplay/__init__.py'
```

```{python}
#| colab: {base_uri: https://localhost:8080/}
#| execution: {iopub.execute_input: '2023-11-30T16:35:09.818095Z', iopub.status.busy: '2023-11-30T16:35:09.817884Z', iopub.status.idle: '2023-11-30T16:35:10.193029Z', shell.execute_reply: '2023-11-30T16:35:10.192647Z', shell.execute_reply.started: '2023-11-30T16:35:09.818079Z'}
#| jupyter: {outputs_hidden: false, source_hidden: false}
%load_ext autoreload
%autoreload 2
import os
import sys
import ezpz

os.environ['COLORTERM'] = 'truecolor'
if sys.platform == 'darwin':
    # If running on MacOS:
    # os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
    os.environ['TORCH_DEVICE'] = 'cpu'
# -----------------------------------------------

logger = ezpz.get_logger()

import wordplay
logger.info(wordplay.__file__)
```

## Build Trainer

Explicitly, we:

1. `setup_torch(...)`
2. Build `cfg: DictConfig = get_config(...)`
3. Instnatiate `config: ExperimentConfig = instantiate(cfg)`
4. Build `trainer = Trainer(config)`

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 481}
#| jupyter: {outputs_hidden: false, source_hidden: false}
import os
import numpy as np
from ezpz import setup
from hydra.utils import instantiate
from wordplay.configs import get_config, PROJECT_ROOT
from wordplay.trainer import Trainer

HF_DATASETS_CACHE = PROJECT_ROOT.joinpath('.cache', 'huggingface')
HF_DATASETS_CACHE.mkdir(exist_ok=True, parents=True)

os.environ['HF_DATASETS_CACHE'] = HF_DATASETS_CACHE.as_posix()

BACKEND = 'DDP'

rank = setup(
    framework='pytorch',
    backend=BACKEND,
    seed=1234,
)

cfg = get_config(
    [
        'data=shakespeare',
        'model=shakespeare',
        'model.batch_size=1',
        'model.block_size=128',
        'optimizer=shakespeare',
        'train=shakespeare',
        f'train.backend={BACKEND}',
        'train.compile=false',
        'train.dtype=bfloat16',
        'train.max_iters=500',
        'train.log_interval=10',
        'train.eval_interval=50',
    ]
)
config = instantiate(cfg)
```

### Build `Trainer` object

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 113}
#| execution: {iopub.execute_input: '2023-11-30T16:35:10.194033Z', iopub.status.busy: '2023-11-30T16:35:10.193752Z', iopub.status.idle: '2023-11-30T16:36:28.944914Z', shell.execute_reply: '2023-11-30T16:36:28.943625Z', shell.execute_reply.started: '2023-11-30T16:35:10.194017Z'}
#| jupyter: {outputs_hidden: false, source_hidden: false}
trainer = Trainer(config)
```

## Prompt (**prior** to training)

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 97}
#| jupyter: {outputs_hidden: false, source_hidden: false}
query = "What is an LLM?"
outputs = trainer.evaluate(
    query,
    num_samples=1,
    max_new_tokens=256,
    top_k=16,
    display=False
)
logger.info(f"['prompt']: '{query}'")
logger.info("['response']:\n\n" + fr"{outputs['0']['raw']}")
```

## Train Model

|  name  |       description            |
|:------:|:----------------------------:|
| `step` | Current training step        |
| `loss` | Loss value                   |
| `dt`   | Time per step (in **ms**)    |
| `sps`  | Samples per second           |
| `mtps` | (million) Tokens per sec     |
| `mfu`  | Model Flops utilization[^1]  |
^legend: #tbl-legend

[^1]: in units of A100 `bfloat16` peak FLOPS

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 35}
#| jupyter: {outputs_hidden: false, source_hidden: false}
trainer.config.device_type
```

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 449}
from rich import print

print(trainer.model)
```

## (partial) Training:

We'll first train for 500 iterations and then evaluate the models performance on the same prompt:

> What is an LLM?

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 625}
#| execution: {iopub.status.busy: '2023-11-30T16:36:28.946773Z', iopub.status.idle: '2023-11-30T16:36:28.946965Z', shell.execute_reply: '2023-11-30T16:36:28.946874Z', shell.execute_reply.started: '2023-11-30T16:36:28.946865Z'}
#| jupyter: {outputs_hidden: false, source_hidden: false}
trainer.train(train_iters=500)
```

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 321}
import time

query = "What is an LLM?"
t0 = time.perf_counter()
outputs = trainer.evaluate(
    query,
    num_samples=1,
    max_new_tokens=256,
    top_k=16,
    display=False
)
logger.info(f'took: {time.perf_counter() - t0:.4f}s')
logger.info(f"['prompt']: '{query}'")
logger.info("['response']:\n\n" + fr"{outputs['0']['raw']}")
```

## Resume Training...

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 832}
trainer.train()
```

## Evaluate Model

```{python}
#| colab: {base_uri: https://localhost:8080/, height: 209}
#| jupyter: {outputs_hidden: false, source_hidden: false}
import time

query = "What is an LLM?"
t0 = time.perf_counter()
outputs = trainer.evaluate(
    query,
    num_samples=1,
    max_new_tokens=256,
    top_k=2,
    display=False
)
logger.info(f'took: {time.perf_counter() - t0:.4f}s')
logger.info(f"['prompt']: '{query}'")
logger.info("['response']:\n\n" + fr"{outputs['0']['raw']}")
```


