@samitouri / QOSami-colab-cli
Code Commits Issues Pulls Sessions CI/CD
main
main
py 153 lines 5.23 KB
Raw
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()}")