docs: add demo video, platform disclaimer, MCP link, and finetune example (#44)
- README: embed CLI demo (docs/demo.webm, 929 KB VP9 WebM) at the top via HTML5 <video> tag with a raw URL pointing at main - README: add Linux/macOS-only platform support disclaimer - README: add pointer to googlecolab/colab-mcp for in-notebook agent workflows - examples/finetune_run.py: new QLoRA fine-tuning example for Gemma 3 1B on a T4 via 'colab exec -f', demonstrating end-to-end script execution
Spencer Shumway committed
Jun 4, 2026 at 14:19 UTC
a9d0277ca65d8a9db6c86450ab18b2da7a422b25
3 files changed
+161
README.md
+8
@@ -4,6 +4,14 @@ A command-line interface for Google Colab. Provision high-performance CPU, GPU,
4
5
Designed to support seamless developer productivity, headless automation, and AI agent integrations.
6
7
+<video src="https://github.com/googlecolab/google-colab-cli/raw/main/docs/demo.webm" controls width="100%"></video>
8
+
9
+> [!NOTE]
10
+> **Platform support:** the Colab CLI currently supports **Linux and macOS** only. Windows is not supported at this time.
11
+
12
+> [!TIP]
13
+> Looking for in-notebook, interactive agent-assisted coding instead of a terminal workflow? See the [Colab MCP Server](https://github.com/googlecolab/colab-mcp).
14
+
15
---
16
17
## Key Features
docs/demo.webm
Binary files /dev/null and b/docs/demo.webm differ
examples/finetune_run.py
new
+153
@@ -0,0 +1,153 @@
1
+# -----------------------------------------------------------------------------
2
+# Setup: HF_TOKEN is required
3
+# -----------------------------------------------------------------------------
4
+# This script downloads google/gemma-3-1b-it, which is a gated model. You need
5
+# a Hugging Face access token. Get one at https://huggingface.co/settings/tokens
6
+# and accept the model license at https://huggingface.co/google/gemma-3-1b-it
7
+#
8
+# 1. Set HF_TOKEN in your local shell, persistently:
9
+#
10
+# echo 'export HF_TOKEN=hf_yourTokenHere' | tee -a ~/.zshrc ~/.bashrc
11
+# source ~/.zshrc # or open a new terminal
12
+# echo $HF_TOKEN # verify — should print your token
13
+#
14
+# 2. Pipe the local env var into the colab kernel before running this script:
15
+#
16
+# echo "import os; os.environ['HF_TOKEN'] = '$HF_TOKEN'" | colab exec
17
+#
18
+# 3. Verify the kernel received it:
19
+#
20
+# echo 'import os; print(bool(os.environ.get("HF_TOKEN")))' | colab exec
21
+# # → should print: True
22
+#
23
+# 4. Run this script:
24
+#
25
+# colab exec -f finetune_run.py
26
+#
27
+# Note: HF_TOKEN lives in the colab kernel for the lifetime of the session.
28
+# If you `colab stop` or the session expires, you'll need to re-pipe it (step 2).
29
+# -----------------------------------------------------------------------------
30
+
31
+import os
32
+
33
+os.system("pip install -q -U 'bitsandbytes>=0.46.1'")
34
+
35
+import torch
36
+from datasets import load_dataset
37
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
38
+from peft import LoraConfig, get_peft_model
39
+from trl import SFTConfig, SFTTrainer
40
+
41
+MODEL_ID = "google/gemma-3-1b-it"
42
+NUM_SAMPLES = 200 # demo size; bump to 5000+ for a real run
43
+MAX_STEPS = 60 # demo cap; set to -1 for full-epoch training
44
+
45
+# -------- Data --------
46
+# philschmid/gretel-synthetic-text-to-sql has sql_prompt, sql_context, sql.
47
+# We hand SFTTrainer a "messages" column and let it apply the chat template.
48
+print("Loading dataset...")
49
+dataset = load_dataset("philschmid/gretel-synthetic-text-to-sql", split="train").select(
50
+ range(NUM_SAMPLES)
51
+)
52
+
53
+
54
+def to_messages(example):
55
+ user_msg = (
56
+ "You are a SQL expert. Given the schema, write a SQL query that "
57
+ "answers the question. Reply with only the SQL.\n\n"
58
+ f"Schema:\n{example['sql_context']}\n\n"
59
+ f"Question:\n{example['sql_prompt']}"
60
+ )
61
+ return {
62
+ "messages": [
63
+ {"role": "user", "content": user_msg},
64
+ {"role": "assistant", "content": example["sql"]},
65
+ ]
66
+ }
67
+
68
+
69
+dataset = dataset.map(to_messages, remove_columns=dataset.column_names)
70
+
71
+# -------- Model (4-bit QLoRA, bf16 throughout) --------
72
+# Everything is bf16 — matches Gemma's natural dtype, matches TRL's default
73
+# T4 (Turing) has no hardware bf16, so this is slower than fp16 would be (~2x)
74
+print(f"Loading {MODEL_ID} in 4-bit...")
75
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
76
+
77
+model = AutoModelForCausalLM.from_pretrained(
78
+ MODEL_ID,
79
+ quantization_config=BitsAndBytesConfig(
80
+ load_in_4bit=True,
81
+ bnb_4bit_quant_type="nf4",
82
+ bnb_4bit_compute_dtype=torch.bfloat16,
83
+ ),
84
+ device_map="auto",
85
+)
86
+
87
+model = get_peft_model(
88
+ model,
89
+ LoraConfig(
90
+ r=16,
91
+ lora_alpha=32,
92
+ target_modules="all-linear",
93
+ task_type="CAUSAL_LM",
94
+ ),
95
+)
96
+# Required for QLoRA backward: makes the embedding output require grad so that
97
+# gradients can flow into the LoRA params attached to layers downstream of the
98
+# frozen 4-bit base.
99
+model.enable_input_require_grads()
100
+model.print_trainable_parameters()
101
+
102
+# -------- Train --------
103
+# All other knobs use SFTConfig defaults (which include bf16=True,
104
+# gradient_checkpointing=True, logging_steps=10). The overrides below are just
105
+# the demo cap, batch sizing that fits T4 VRAM, and silencing wandb/tensorboard.
106
+print("Training...")
107
+trainer = SFTTrainer(
108
+ model=model,
109
+ train_dataset=dataset,
110
+ processing_class=tokenizer,
111
+ args=SFTConfig(
112
+ output_dir="./results",
113
+ max_steps=MAX_STEPS,
114
+ per_device_train_batch_size=2,
115
+ gradient_accumulation_steps=2,
116
+ # Standard QLoRA LR. SFTConfig defaults to 2e-5, which is too low for
117
+ # LoRA adapters to learn anything meaningful in 60 steps.
118
+ learning_rate=2e-4,
119
+ # Compute loss only on the assistant's SQL, not on the schema/question.
120
+ assistant_only_loss=True,
121
+ # Off so KV cache works during the inference step at the end.
122
+ gradient_checkpointing=False,
123
+ report_to="none",
124
+ ),
125
+)
126
+trainer.train()
127
+
128
+# -------- Save --------
129
+out_dir = "./gemma-3-1b-qlora-adapter"
130
+trainer.model.save_pretrained(out_dir)
131
+tokenizer.save_pretrained(out_dir)
132
+print(f"Saved adapter to {out_dir}")
133
+
134
+# -------- Inference check --------
135
+sample = dataset[0]
136
+prompt = tokenizer.apply_chat_template(
137
+ sample["messages"][:1], # just the user turn
138
+ tokenize=False,
139
+ add_generation_prompt=True,
140
+)
141
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
142
+with torch.no_grad():
143
+ out_ids = model.generate(
144
+ **inputs,
145
+ max_new_tokens=256,
146
+ do_sample=False,
147
+ pad_token_id=tokenizer.pad_token_id,
148
+ )
149
+generated = tokenizer.decode(
150
+ out_ids[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
151
+)
152
+print(f"\nGold: {sample['messages'][1]['content']}")
153
+print(f"Model: {generated.strip()}")