Skip to content

Commit 52e6b7f

Browse files
committed
added in line editable commands for env variables in tutorials
1 parent 4962bb8 commit 52e6b7f

7 files changed

Lines changed: 184 additions & 11 deletions

File tree

docs/_static/css/custom.css

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,3 +209,35 @@ html[data-theme="dark"] .button-primary:visited:hover {
209209
padding-inline: 50px;
210210
padding-bottom: 24px;
211211
}
212+
213+
/* Inline inputs for editable commands */
214+
.inline-input {
215+
background-color: rgba(255, 255, 255, 0.1);
216+
color: inherit;
217+
border: 1px solid #555;
218+
border-radius: 3px;
219+
padding: 2px 6px;
220+
font-family: inherit;
221+
font-size: inherit;
222+
display: inline-block;
223+
vertical-align: middle;
224+
margin: 0 2px;
225+
box-sizing: content-box;
226+
}
227+
228+
.inline-input:focus {
229+
outline: none;
230+
border-color: #1A73E8;
231+
background-color: rgba(255, 255, 255, 0.2);
232+
}
233+
234+
html[data-theme="light"] .inline-input {
235+
background-color: rgba(0, 0, 0, 0.05);
236+
border-color: #ccc;
237+
color: #333;
238+
}
239+
240+
html[data-theme="light"] .inline-input:focus {
241+
background-color: rgba(0, 0, 0, 0.1);
242+
border-color: #1A73E8;
243+
}
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
/**
2+
* Handles inline editable commands in documentation.
3+
* Replaces placeholders in code blocks with inline input fields.
4+
*/
5+
document.addEventListener('DOMContentLoaded', () => {
6+
const codeBlocks = document.querySelectorAll('div.highlight-sh pre, div.highlight-bash pre, div.highlight-default pre');
7+
8+
codeBlocks.forEach(block => {
9+
10+
const originalHTML = block.innerHTML;
11+
12+
const placeholders = [
13+
"<your virtual env name>",
14+
"<model name>",
15+
"<tokenizer path>",
16+
"<Hugging Face access token>",
17+
"<output directory to store run logs>",
18+
"<name for this run>",
19+
"<number of fine-tuning steps to run>",
20+
"<batch size per device>",
21+
"<Hugging Face dataset name>",
22+
"<data split for train>",
23+
"<data columns to train on>",
24+
"<gcs path for MaxText checkpoint>",
25+
"<Google Cloud Project ID>",
26+
"<Name of GKE Cluster>",
27+
"<GKE Cluster Zone>",
28+
"<Name of Workload>",
29+
"<TPU Type>",
30+
"<GCS Path for Output/Logs>",
31+
"<Fine-Tuning Steps>",
32+
"<Hugging Face Access Token>",
33+
"<Model Name>",
34+
"<Model Tokenizer>",
35+
"<Hugging Face Dataset Name>",
36+
"<Data Split for Train>",
37+
"<Data Columns to Train on>",
38+
"<cluster name>",
39+
"<GCP project ID>",
40+
"<zone name>",
41+
"<path/to/gcr.io>",
42+
"<number of slices>",
43+
"<Flag to use zarr3>",
44+
"<Flag to use ocdbt>",
45+
"<Hugging Face Model>",
46+
"<MaxText Model>",
47+
"<Tokenizer>",
48+
"<Name for this run>",
49+
"<Docker Image Name>"
50+
];
51+
52+
let newHTML = originalHTML;
53+
54+
placeholders.forEach(placeholder => {
55+
// 1. create robust regex for this placeholder
56+
// escape chars
57+
const escapeRegex = (string) => string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
58+
59+
const htmlEscapedKey = placeholder
60+
.replace(/&/g, '&amp;')
61+
.replace(/</g, '&lt;')
62+
.replace(/>/g, '&gt;');
63+
64+
let pattern = '';
65+
for (let i = 0; i < htmlEscapedKey.length; i++) {
66+
const char = htmlEscapedKey[i];
67+
pattern += escapeRegex(char) + '(?:<[^>]+>)*';
68+
}
69+
70+
const regex = new RegExp(pattern, 'g');
71+
72+
// Replace with an input element
73+
// We use the original placeholder text as placeholder for the input
74+
const inputHTML = `<input class="inline-input" placeholder="${placeholder}" style="width: ${placeholder.length + 2}ch;" />`;
75+
76+
newHTML = newHTML.replace(regex, inputHTML);
77+
});
78+
79+
if (newHTML !== originalHTML) {
80+
block.innerHTML = newHTML;
81+
}
82+
});
83+
84+
// Add event listeners to newly created inputs to auto-resize
85+
document.querySelectorAll('.inline-input').forEach(input => {
86+
input.addEventListener('input', function () {
87+
this.style.width = Math.max(this.value.length, this.placeholder.length) + 2 + 'ch';
88+
});
89+
});
90+
91+
/**
92+
* Intercept copy button clicks to include user input values.
93+
* Runs in capture phase to precede sphinx-copybutton's listener.
94+
*/
95+
document.addEventListener('click', (event) => {
96+
// Check if the clicked element is a copy button or inside one
97+
const button = event.target.closest('.copybtn');
98+
if (!button) return;
99+
100+
// Find the associated code block
101+
// Sphinx-copybutton places the button inside .highlight usually
102+
const highlightDiv = button.closest('.highlight');
103+
if (!highlightDiv) return;
104+
105+
const inputs = highlightDiv.querySelectorAll('input.inline-input');
106+
if (inputs.length === 0) return;
107+
108+
const swaps = [];
109+
inputs.forEach(input => {
110+
// Create a temporary span with the input's current value
111+
const span = document.createElement('span');
112+
// If value is empty, fallback to placeholder to match original text behavior
113+
const val = input.value;
114+
span.textContent = val ? val : input.placeholder;
115+
116+
// Mimic input appearance slightly if needed, but plain text is what we want copied
117+
span.style.color = val ? 'inherit' : 'gray';
118+
119+
input.replaceWith(span);
120+
swaps.push({ input, span });
121+
});
122+
123+
// Revert immediately after the current event loop
124+
setTimeout(() => {
125+
swaps.forEach(({ input, span }) => {
126+
span.replaceWith(input);
127+
});
128+
}, 0);
129+
}, true);
130+
});

docs/conf.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
html_theme = "sphinx_book_theme"
4949
html_static_path = ["_static"]
5050
html_css_files = ["css/custom.css"]
51+
html_js_files = ["js/editable_commands.js"]
5152
html_logo = "_static/maxtext.png"
5253

5354
# -- Options for myst ----------------------------------------------

docs/tutorials/posttraining/rl.md

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,9 @@ Setup following environment variables before running GRPO/GSPO:
6969

7070
```bash
7171
# -- Model configuration --
72-
export HF_MODEL='llama3.1-8b-Instruct'
73-
export MODEL='llama3.1-8b'
74-
export TOKENIZER='meta-llama/Llama-3.1-8B-Instruct'
72+
export HF_MODEL=<Hugging Face Model> # e.g. 'llama3.1-8b-Instruct'
73+
export MODEL=<MaxText Model> # e.g. 'llama3.1-8b'
74+
export TOKENIZER=<Tokenizer> # e.g. 'meta-llama/Llama-3.1-8B-Instruct'
7575
export HF_TOKEN=<Hugging Face access token>
7676

7777
# -- MaxText configuration --

docs/tutorials/posttraining/rl_on_multi_host.md

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -39,26 +39,36 @@ Setup following environment variables:
3939

4040
```bash
4141
# -- Model configuration --
42-
export HF_MODEL='llama3.1-70b-Instruct'
43-
export MODEL='llama3.1-70b'
44-
export TOKENIZER='meta-llama/Llama-3.1-70B-Instruct'
42+
export HF_MODEL=<Hugging Face Model> # e.g. 'llama3.1-70b-Instruct'
43+
export MODEL=<MaxText Model> # e.g. 'llama3.1-70b'
44+
export TOKENIZER=<Tokenizer> # e.g. 'meta-llama/Llama-3.1-70B-Instruct'
4545
export HF_TOKEN=<Hugging Face access token>
4646

4747
# -- MaxText configuration --
4848
export BASE_OUTPUT_DIRECTORY=<output directory to store run logs> # e.g., gs://my-bucket/my-output-directory
49-
export RUN_NAME=llama-3-70b-grpo
49+
export RUN_NAME=<Name for this run> # e.g., llama-3-70b-grpo
5050
export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/${RUN_NAME}/0/items
5151

5252
# -- Workload configuration --
5353
export WORKLOAD=${RUN_NAME}
54-
export TPU_TYPE='v5p-128'
54+
export TPU_TYPE=<TPU Type> # e.g., 'v5p-128'
5555
export TPU_CLUSTER=<cluster name>
5656
export PROJECT_ID=<GCP project ID>
5757
export ZONE=<zone name>
5858
```
5959

6060
## Get your model checkpoint
6161

62+
### Option 1: Using an existing MaxText checkpoint
63+
64+
If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section.
65+
66+
```bash
67+
export MAXTEXT_CKPT_PATH=<gcs path for MaxText checkpoint> # e.g., gs://my-bucket/my-model-checkpoint/0/items
68+
```
69+
70+
### Option 2: Converting from a Hugging Face checkpoint
71+
6272
You can convert a Hugging Face checkpoint to MaxText format using the `src/MaxText/utils/ckpt_conversion/to_maxtext.py` script. This is useful if you have a pre-trained model from Hugging Face that you want to use with MaxText.
6373

6474
First, ensure you have the necessary dependencies installed. Then, run the conversion script on a CPU machine. For large models, it is recommended to use the `--lazy_load_tensors` flag to reduce memory usage during conversion. \

docs/tutorials/posttraining/sft.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ We use [Tunix](https://github.com/google/tunix), a JAX-based library designed fo
2424
In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get started!
2525

2626
## Install dependencies
27+
2728
```sh
2829
# 1. Clone the repository
2930
git clone https://github.com/AI-Hypercomputer/maxtext.git

docs/tutorials/posttraining/sft_on_multi_host.md

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,7 @@ bash dependencies/scripts/docker_build_dependency_image.sh WORKFLOW=post-trainin
5151
### 1.3. Upload the Docker image to Artifact Registry
5252
> **Note:** You will need the [**Artifact Registry Writer**](https://docs.cloud.google.com/artifact-registry/docs/access-control#permissions) role to push Docker images to your project's Artifact Registry and to allow the cluster to pull them during workload execution. If you don't have this permission, contact your project administrator to grant you this role through "Google Cloud Console -> IAM -> Grant access".
5353
```bash
54-
# Replace `$USER_runner` with your desired image name
55-
export DOCKER_IMAGE_NAME=${USER}_runner
54+
export DOCKER_IMAGE_NAME=<Docker Image Name>
5655
bash dependencies/scripts/docker_upload_runner.sh CLOUD_IMAGE_NAME=$DOCKER_IMAGE_NAME
5756
```
5857
The `docker_upload_runner.sh` script uploads your Docker image to Artifact Registry.
@@ -73,7 +72,7 @@ export ZONE=<GKE Cluster Zone>
7372
# -- Workload Configuration --
7473
export WORKLOAD_NAME=<Name of Workload> # e.g., sft-$(date +%s)
7574
export TPU_TYPE=<TPU Type> # e.g., v6e-256
76-
export TPU_SLICE=1
75+
export TPU_SLICE=<number of slices>
7776
export DOCKER_IMAGE="gcr.io/${PROJECT}/${DOCKER_IMAGE_NAME}"
7877

7978
# -- MaxText Configuration --

0 commit comments

Comments
 (0)