mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
44 KiB
44 KiB
In [1]:
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.Cloud AI Alphagenome Fine tuning Notebook
 Open in Colab Enterprise |
In [ ]:
from IPython.display import clear_output
! PIP_NO_BINARY=pyBigWig pip install git+https://github.com/google-deepmind/alphagenome_research.git
clear_output()In [ ]:
import jax
# We need >0.9.0 jax libs.
# Check the jax version.
import jaxlib
print(f"{jax.__version__=}")
print(f"{jaxlib.__version__=}")In [ ]:
# Uninstall the previous version.
# Run only if version < 0.9.0.
# Restart runtime/kernel after uninstalling.
# Run from next cell after the kernel restart
if jax.__version__ != "0.9.0":
print(f"Unistalling {jax.__version__}")
! pip uninstall -y jax jaxlib jax_cuda12_pluginIn [ ]:
# Install specific vesion.
# Run only once after unistalling the jax packages.
# Restart the runtime/kernel.
# Run from next cell after the kernel restart.
!pip install --upgrade jax[cuda12_pip]==0.9.0 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlIn [ ]:
# Install specific vesion.
# Run only once after the upgrade.
# Restart the runtime/kernel.
# Run from next cell after the kernel restart.
!pip install jax_cuda12_plugin==0.9.0In [ ]:
import dataclasses
import os
import pprint
import subprocess
import time
import haiku as hk
import huggingface_hub
import numpy as np
import optax
import orbax.checkpoint as ocp
import pandas as pd
from alphagenome.data import fold_intervals, genome
from alphagenome.visualization import plot_components
from alphagenome_research.finetuning import dataset as dataset_lib
from alphagenome_research.finetuning import finetune
from alphagenome_research.model import dna_model
from alphagenome_research.model.metadata import metadata as metadata_lib
from etils import epath
from google.cloud import secretmanager
from huggingface_hub import login
from jax import errors
from jax.experimental import mesh_utils
from jax.sharding import Mesh
from jax.sharding import PartitionSpec as PIn [ ]:
import tensorflow as tf
# Hide local GPUs/TPUs. TensorFlow only used for data loading.
tf.config.set_visible_devices([], "GPU")
tf.config.set_visible_devices([], "TPU")In [ ]:
import jax
# we need >0.9.0 jax libs
# check the jax versions
import jaxlib
print(f"{jax.__version__=}")
print(f"{jaxlib.__version__=}")In [ ]:
# Ensure you are running an Enterprise Colab with a runtime that
# have GPUs
def gpu_info() -> None:
"""Prints the GPU information."""
try:
backend = jax.default_backend()
if backend == "gpu":
num_gpus = jax.local_device_count()
print(f"JAX is using GPU backend with {num_gpus} GPU(s).")
else:
print(f"JAX default backend is {backend}, not GPU.")
except errors.JaxRuntimeError as e:
print(f"JAX runtime error occurred while detecting devices: {e}")
gpu_info()In [ ]:
# --- Google Cloud Platform variables ---
PROJECT_ID = "<your project>" # @param {type:'string'}
BUCKET_NAME = "<your bucket>" # @param {type: 'string'}
# ---- Local file directories ----
LOCAL_TAR_DIR = "/tmp/tar_outputs"
LOCAL_BIGWIG_DIR = "/tmp/bigwig"
LOCAL_FASTA_DIR = "/tmp/fasta"
LOCAL_HUMAN_SEQ_DIR = "/tmp/example_regions_path"
SAVE_CHECKPOINT_DIR = "/tmp/checkpoint"
# --- AlphaGenome finetuning Model params---
LEARNING_RATE = 5e-4
NUM_TRAIN_STEPS = 1000
MODEL_VERSION = dna_model.ModelVersion.FOLD_0
SEQUENCE_LENGTH = int(2**20)
BATCH_SIZE = 1 # Per device
ORGANISM = dna_model.Organism.HOMO_SAPIENS
# --- derived variables ---
from datetime import datetime
GCS_AGFT_NAME = f"ag-ft-{BATCH_SIZE}-{NUM_TRAIN_STEPS}"
GCS_AGFT_NAME = f"{GCS_AGFT_NAME}-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
GCS_AGFT_NAME = f"{GCS_AGFT_NAME}.tar.gz"
GCS_AGFT_NAME = GCS_AGFT_NAME.replace("_", "-")
GCS_AGFT_PATH = f"{BUCKET_NAME}/finetune/{GCS_AGFT_NAME}"
GCS_AGFT_PLOT_NAME = GCS_AGFT_NAME.replace(".tar.gz", ".png")
GCS_AGFT_PLOT_PATH = f"{BUCKET_NAME}/finetune/{GCS_AGFT_PLOT_NAME}"In [ ]:
# Check if the file does NOT exist
if not os.path.exists(LOCAL_FASTA_DIR):
! mkdir -p $LOCAL_FASTA_DIR
! echo wget -P "$LOCAL_FASTA_DIR" https://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_46/GRCh38.p14.genome.fa.gz
! wget -P "$LOCAL_FASTA_DIR" https://ftp.ebi.ac.uk/pub/databases/gencode/Gencode_human/release_46/GRCh38.p14.genome.fa.gz
# unzip, create an index file, .tar.gz, and upload to gcs
! echo gunzip $LOCAL_FASTA_DIR/GRCh38.p14.genome.fa.gz
! gunzip $LOCAL_FASTA_DIR/GRCh38.p14.genome.fa.gz
# install samtools to create the inded file
! echo apt install samtools
! apt install samtools
# create the index file
# the tool creates GRCh38.p14.genome.fa.fai and stores in the same dir.
! echo samtools faidx $LOCAL_FASTA_DIR/GRCh38.p14.genome.fa
! samtools faidx $LOCAL_FASTA_DIR/GRCh38.p14.genome.fa
print("All set to use fasta files in the training.")
else:
print("Going to use already prepared fasta files.")In [ ]:
if not os.path.exists(LOCAL_HUMAN_SEQ_DIR):
# create dir
! echo mkdir $LOCAL_HUMAN_SEQ_DIR
# download the file
! echo wget -P "$LOCAL_HUMAN_SEQ_DIR" https://github.com/calico/borzoi/raw/5c9358222b5026abb733ed5fb84f3f6c77239b37/data/sequences_human.bed.gz
! wget -P "$LOCAL_HUMAN_SEQ_DIR" https://github.com/calico/borzoi/raw/5c9358222b5026abb733ed5fb84f3f6c77239b37/data/sequences_human.bed.gz
# unzip the file
! echo gunzip $LOCAL_HUMAN_SEQ_DIR/sequences_human.bed.gz
! gunzip $LOCAL_HUMAN_SEQ_DIR/sequences_human.bed.gz
print("All set to use human sequence file in the training.")
else:
print("Going to use already prepared human sequence files.")In [ ]:
if not os.path.exists(LOCAL_BIGWIG_DIR):
# create the temp dire
! mkdir -p $LOCAL_BIGWIG_DIR
# downloads the big wig files
! pushd $LOCAL_BIGWIG_DIR && curl \
-C - \
-Z -O https://storage.googleapis.com/alphagenome/reference/encode/hg38/ENCFF018EZY.bigWig \
-O https://storage.googleapis.com/alphagenome/reference/encode/hg38/ENCFF904TSK.bigWig \
-O https://storage.googleapis.com/alphagenome/reference/encode/hg38/ENCFF218CLQ.bigWig && popd
print("All set to use the bigwig files in the training.")
else:
print("Going to use already prepared bigwig files.")In [ ]:
TRACK_METADATA = pd.DataFrame(
data=[
[
"RNA_SEQ",
"UBERON:0000948 total RNA-seq",
"+",
f"{LOCAL_BIGWIG_DIR}/ENCFF018EZY.bigWig",
],
[
"RNA_SEQ",
"UBERON:0000948 total RNA-seq",
"-",
f"{LOCAL_BIGWIG_DIR}/ENCFF904TSK.bigWig",
],
[
"DNASE",
"EFO:0005337 DNase-seq",
".",
f"{LOCAL_BIGWIG_DIR}/ENCFF218CLQ.bigWig",
],
],
columns=["output_type", "name", "strand", "file_path"],
)
TRACK_METADATAIn [ ]:
def build_output_metadata(
track_metadata: pd.DataFrame,
) -> metadata_lib.AlphaGenomeOutputMetadata:
"""Builds AlphaGenomeOutputMetadata from the track metadata DataFrame.
Args:
track_metadata: A pandas DataFrame containing metadata for the tracks,
including 'output_type', 'name', 'strand', and 'file_path'.
Returns:
A dict mapping organism to AlphaGenomeOutputMetadata.
"""
required_cols = {"file_path", "name", "output_type", "strand"}
if not required_cols.issubset(track_metadata.columns):
raise ValueError(
f"track_metadata must have columns {required_cols}. Missing: {required_cols - set(track_metadata.columns)}."
)
metadata = {}
for output_type, df_group in track_metadata.groupby("output_type"):
try:
output_type = dna_model.OutputType[str(output_type)]
except KeyError as e:
raise ValueError(f"Unknown output_type: {output_type}") from e
metadata[output_type.name.lower()] = df_group
return metadata_lib.AlphaGenomeOutputMetadata(**metadata)
output_metadata = {
dna_model.Organism.HOMO_SAPIENS: build_output_metadata(TRACK_METADATA)
}In [ ]:
ds_iter = finetune.get_dataset_iterator(
batch_size=BATCH_SIZE * jax.local_device_count(),
sequence_length=SEQUENCE_LENGTH,
output_metadata=output_metadata[ORGANISM],
organism=ORGANISM,
model_version=MODEL_VERSION,
subset=fold_intervals.Subset.TRAIN,
fasta_path=f"{LOCAL_FASTA_DIR}/GRCh38.p14.genome.fa",
example_regions_path=f"{LOCAL_HUMAN_SEQ_DIR}/sequences_human.bed",
)In [ ]:
# validate the shape
batch = next(ds_iter)
pprint.pprint(jax.tree.map(np.shape, batch))In [ ]:
# setup Huggingface credential and download the base model weights
def setup_huggingface_auth(secret_id="HUGGINGFACE_API_TOKEN", version_id="latest"):
"""Fetches HF token from Secret Manager and configures auth."""
try:
project_id = os.environ.get("GOOGLE_CLOUD_PROJECT")
print(f"{project_id=}")
if not project_id:
try:
project_id = (
subprocess.check_output(
["gcloud", "config", "get-value", "project"]
)
.decode("utf-8")
.strip()
)
except subprocess.CalledProcessError:
print("Could not automatically determine GCP Project ID.")
return None
client = secretmanager.SecretManagerServiceClient()
name = f"projects/{project_id}/secrets/{secret_id}/versions/{version_id}"
response = client.access_secret_version(name=name)
hf_token = response.payload.data.decode("UTF-8").strip()
if not hf_token:
print(f"Secret {secret_id} is empty.")
return None
print("Hugging Face token retrieved from Secret Manager.")
return hf_token
except Exception as e:
print(f"Error setting up Hugging Face auth: {e}")
return None
hf_token = setup_huggingface_auth()
if hf_token:
# Option 1: Log in using huggingface-cli
# This makes the token available for CLI commands and many libraries.
try:
# Use subprocess to handle the interactive nature of login
process = subprocess.Popen(
["huggingface-cli", "login", "--token", hf_token],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
stdout, stderr = process.communicate()
if process.returncode == 0:
print("Hugging Face CLI login successful.")
else:
print(f"Hugging Face CLI login failed: {stderr.decode()}")
print(stdout.decode())
except FileNotFoundError:
print("huggingface-cli not found. Make sure huggingface_hub is installed.")
except Exception as e:
print(f"Error during huggingface-cli login: {e}")
# Option 2: Set as environment variable (useful for some tools)
os.environ["HF_TOKEN"] = hf_token
print("HF_TOKEN environment variable set.")
# Option 3: Programmatic login with huggingface_hub
try:
login(token=hf_token)
print("huggingface_hub programmatic login successful.")
except Exception as e:
print(f"huggingface_hub login error: {e}")
# Now try accessing the gated model again, for example:
# from huggingface_hub import hf_hub_download
repo = f"google/alphagenome-{MODEL_VERSION.name.lower().replace('_', '-')}"
checkpoint_path = huggingface_hub.snapshot_download(repo_id=repo)
checkpointer = ocp.StandardCheckpointer()
params_base, state_base = checkpointer.restore(checkpoint_path)In [ ]:
num_devices = jax.local_device_count()
devices = mesh_utils.create_device_mesh((num_devices,))
mesh = Mesh(devices, axis_names=("data",))
data_sharding = P("data")
replicated_sharding = P()In [ ]:
forward_fn = finetune.get_forward_fn(output_metadata)
with jax.set_mesh(mesh):
batch = jax.device_put(batch, data_sharding)
params_ft, state_ft = jax.jit(
forward_fn.init,
in_shardings=(replicated_sharding, data_sharding),
out_shardings=replicated_sharding,
)(jax.random.PRNGKey(0), batch)In [ ]:
params_ft_head = hk.data_structures.filter(
lambda module_name, *_: "head" in module_name, params_ft
)
params_base_no_head = hk.data_structures.filter(
lambda module_name, *_: "head" not in module_name, params_base
)
params = hk.data_structures.merge(params_base_no_head, params_ft_head)
state = state_base
optimizer = optax.chain(
optax.clip_by_global_norm(0.5),
optax.adam(LEARNING_RATE),
)
opt_state = optimizer.init(params)
train_step = jax.jit(
finetune.get_train_step(forward_fn.apply, optimizer),
in_shardings=(
replicated_sharding,
replicated_sharding,
replicated_sharding,
data_sharding,
),
out_shardings=(
replicated_sharding,
replicated_sharding,
replicated_sharding,
replicated_sharding,
),
)In [ ]:
path_suffix = datetime.now().strftime("%Y%m%d_%H%M%S")
checkpoint_dir = epath.Path(SAVE_CHECKPOINT_DIR) / path_suffix
checkpoint_dir.mkdir(parents=True, exist_ok=True)
checkpoint_dir = str(checkpoint_dir)
print(f"We will be saving the trained checkpoint at {checkpoint_dir}")In [ ]:
checkpointer = ocp.StandardCheckpointer()
def save(weights, idx):
ckpt_path = os.path.join(checkpoint_dir, "checkpoint_{:05d}".format(idx))
print(f"Saving checkpoint to {ckpt_path}")
checkpointer.save(ckpt_path, weights)
checkpointer.wait_until_finished()
print(f"Saved checkpoint to {ckpt_path}")
return ckpt_pathIn [ ]:
loss, times = [], []
for step in range(NUM_TRAIN_STEPS):
start_time = time.time()
try:
batch = next(ds_iter)
except StopIteration:
print("Dataset exhausted")
break
with jax.set_mesh(mesh):
batch = jax.device_put(batch, data_sharding)
params, state, opt_state, scalars = train_step(params, state, opt_state, batch)
loss.append(scalars["loss"])
times.append(time.time() - start_time)
if step % 10 == 1:
print("loss", step, loss[-1], f"SPS: {1./np.mean(times[1:]):.4f}")
print(f"Total Training time: {np.sum(times[1:]):.4f} seconds")
print(f"Average Training time per step: {np.mean(times[1:]):.4f} seconds")
ckpt_path = save((params, state), step + 1)In [ ]:
# Check if the file or directory does NOT exist
# gcs path to where the chkpt will be uploaded
os.makedirs(LOCAL_TAR_DIR, exist_ok=True)
# tar the files
!echo tar -czvf $LOCAL_TAR_DIR/$GCS_AGFT_NAME $ckpt_path
!tar -czvf $LOCAL_TAR_DIR/$GCS_AGFT_NAME $ckpt_path
# list the tar file
!echo ls -l $LOCAL_TAR_DIR/$GCS_AGFT_NAME
!ls -l $LOCAL_TAR_DIR/$GCS_AGFT_NAME
# upload the tar file to gcs bucket
!echo gsutil cp $LOCAL_TAR_DIR/$GCS_AGFT_NAME $GCS_AGFT_PATH
!gsutil cp $LOCAL_TAR_DIR/$GCS_AGFT_NAME $GCS_AGFT_PATH
# list the tar file gcs bucket
!echo gsutil ls $GCS_AGFT_PATH
!gsutil ls $GCS_AGFT_PATHIn [ ]:
# Load default organism settings but overwrite with fine-tuned output metadata.
default_settings_human = dna_model.default_organism_settings()[
dna_model.Organism.HOMO_SAPIENS
]
settings_human_finetune = dataclasses.replace(
default_settings_human,
metadata=output_metadata[dna_model.Organism.HOMO_SAPIENS],
)
model = dna_model.create(
ckpt_path,
organism_settings={dna_model.Organism.HOMO_SAPIENS: settings_human_finetune},
)In [ ]:
interval = genome.Interval(chromosome="chr21", start=46125238, end=46126738).resize(
SEQUENCE_LENGTH
)In [ ]:
preds = model.predict_interval(
interval,
requested_outputs=[dna_model.OutputType.RNA_SEQ],
ontology_terms=None,
)
predsIn [ ]:
true_tracks = dataset_lib.MultiTrackExtractor(
output_metadata[ORGANISM], sequence_length=SEQUENCE_LENGTH
).extract(interval)In [ ]:
def compact_dict(**kwargs):
return {k: v for k, v in kwargs.items() if v is not None}
def plot(*, interval, predictions, targets=None):
if targets is None:
colors = {"pred": "black"}
else:
colors = {"pred": "black", "true": "red"}
fig = plot_components.plot(
[
plot_components.OverlaidTracks(
tdata=compact_dict(
pred=predictions.rna_seq,
true=(
dataclasses.replace(
predictions.rna_seq,
values=targets["rna_seq"].astype(np.float32),
)
if targets is not None
else None
),
),
colors=colors,
shared_y_scale=True,
),
],
interval=interval.resize(int(2**11)),
)
return fig
pred_vs_groundtruth_fig = plot(
predictions=preds, interval=interval, targets=true_tracks
)In [ ]:
# Save the figure to a file
file_path = (
f"{LOCAL_TAR_DIR}/{GCS_AGFT_PLOT_NAME}" # Choose your desired filename and format
)
pred_vs_groundtruth_fig.savefig(file_path)
print(f"Plot saved to {file_path}")
# upload the image file to gcs bucket
!echo gsutil cp $file_path $GCS_AGFT_PLOT_PATH
!gsutil cp $file_path $GCS_AGFT_PLOT_PATH
# list the tar file gcs bucket
!echo gsutil ls $GCS_AGFT_PLOT_PATH
!gsutil ls $GCS_AGFT_PLOT_PATH