Skills › Scientific Agent Skills › deepchem › transfer_learning.py
transfer_learning.py
219 lines · skills/deepchem/scripts/transfer_learning.py at 5f2bc03
1#!/usr/bin/env python3
2"""Fine-tune explicit Hugging Face or compatible DeepChem GROVER checkpoints.
4HF examples support one fully observed, unweighted binary or regression task.
5GROVER requires a local DeepChem component checkpoint and matching architecture.
6No pretrained weights are implied by a model output directory.
7"""
8import argparse
9import json
10import re
11import sys
12from pathlib import Path
14import deepchem as dc
15import numpy as np
16from _common import load_csv, validate_splits, evaluate_splits
17from graph_neural_network import MOLNET_DATASETS, molnet_loader
19PRETRAINED_MODELS = {
20 'chemberta': {
21 'name': 'ChemBERTa',
22 'description': 'RoBERTa pretrained on 100k ZINC SMILES (model card)',
23 'model_id': 'seyonec/ChemBERTa-zinc-base-v1',
24 },
25 'grover': {
26 'name': 'GROVER',
27 'description': 'Transfer a matching DeepChem GROVER embedding checkpoint',
28 'model_id': None,
29 },
30 'molformer': {
31 'name': 'MoLFormer',
32 'description': 'MoLFormer-XL pretrained on 10% ZINC plus 10% PubChem',
33 'model_id': 'ibm-research/MoLFormer-XL-both-10pct',
34 },
35}
38def transfer_featurizer(model_type):
39 if model_type in ('chemberta', 'molformer'):
40 return dc.feat.DummyFeaturizer()
41 if model_type == 'grover':
42 return dc.feat.GroverFeaturizer(features_generator=dc.feat.CircularFingerprint(size=2048))
43 raise ValueError(f'Unknown model type: {model_type}')
46def load_molnet_dataset(dataset_name, model_type):
47 if dataset_name not in MOLNET_DATASETS:
48 raise ValueError(f'Unknown dataset: {dataset_name}')
49 # Raw alias produces RDKit Mol objects in stable 2.8.0, not tokenizer strings.
50 # Disable balancing/normalization here: the HF loss ignores dataset.w.
51 return molnet_loader(dataset_name)(featurizer=transfer_featurizer(model_type),
52 splitter='scaffold', transformers=[])
55def load_custom_dataset(data_path, target_cols, smiles_col, model_type):
56 dataset = load_csv(data_path, target_cols, smiles_col, transfer_featurizer(model_type))
57 return dc.splits.ScaffoldSplitter().train_valid_test_split(
58 dataset, frac_train=.8, frac_valid=.1, frac_test=.1)
61def validate_hf_data(datasets, task_type, n_tasks):
62 if n_tasks != 1:
63 raise ValueError('HF example supports one task; sparse multitask loss needs a custom adapter')
64 validate_splits(datasets, task_type, n_tasks)
65 for dataset in datasets:
66 if not np.all(dataset.w == 1):
67 raise ValueError('HF 2.8.0 loss ignores weights: labels must be complete and weights equal 1')
68 if not all(isinstance(value, str) for value in dataset.X):
69 raise ValueError('HF tokenizer requires SMILES strings')
72def validate_remote_revision(trust_remote_code, revision):
73 if trust_remote_code and (not isinstance(revision, str) or
74 re.fullmatch(r'[0-9a-fA-F]{40}', revision) is None):
75 raise ValueError('Remote code requires a reviewed immutable 40-character commit SHA in --revision')
78def build_hf_model(model_type, task_type, model_id=None, revision=None,
79 trust_remote_code=False, local_files_only=False):
80 from transformers import AutoModelForSequenceClassification, AutoTokenizer
81 from deepchem.models.torch_models import HuggingFaceModel
82 if model_type == 'molformer' and not trust_remote_code:
83 raise ValueError('MoLFormer needs reviewed remote code; supply --trust-remote-code and a pinned --revision')
84 validate_remote_revision(trust_remote_code, revision)
85 model_id = model_id or PRETRAINED_MODELS[model_type]['model_id']
86 options = dict(trust_remote_code=trust_remote_code, local_files_only=local_files_only)
87 if revision:
88 options['revision'] = revision
89 tokenizer = AutoTokenizer.from_pretrained(model_id, **options)
90 model_options = dict(options)
91 if model_type == 'molformer':
92 # Its forward accepts no token_type_ids; generic fast tokenizers add them.
93 tokenizer.model_input_names = ['input_ids', 'attention_mask']
94 model_options['deterministic_eval'] = True
95 network = AutoModelForSequenceClassification.from_pretrained(
96 model_id, num_labels=2 if task_type == 'classification' else 1,
97 problem_type='single_label_classification' if task_type == 'classification' else 'regression',
98 **model_options)
99 # The base checkpoint supplies the encoder; the prediction head may be new.
100 return HuggingFaceModel(model=network, tokenizer=tokenizer, task=task_type,
101 batch_size=16, learning_rate=2e-5, device='cpu')
104def train_hf(model_type, train, valid, test, task_type='classification',
105 n_tasks=1, n_epochs=10, **options):
106 validate_hf_data((train, valid, test), task_type, n_tasks)
107 model = build_hf_model(model_type, task_type, **options)
108 model.fit(train, nb_epoch=n_epochs)
109 # The stable HuggingFaceModel prediction contract is raw logits.
110 results = evaluate_splits(model, (train, valid, test), task_type,
111 logits=task_type == 'classification')
112 return model, results
115def train_chemberta(train_dataset, valid_dataset, test_dataset,
116 task_type='classification', n_tasks=1, n_epochs=10, **options):
117 return train_hf('chemberta', train_dataset, valid_dataset, test_dataset,
118 task_type, n_tasks, n_epochs, **options)
121def train_molformer(train_dataset, valid_dataset, test_dataset,
122 task_type='classification', n_tasks=1, n_epochs=10, **options):
123 return train_hf('molformer', train_dataset, valid_dataset, test_dataset,
124 task_type, n_tasks, n_epochs, **options)
127def build_grover(task_type, n_tasks, config, checkpoint):
128 """Strictly restore only the encoder, leaving a new supervised head."""
129 import torch
130 from deepchem.models.torch_models import GroverModel
131 if not checkpoint or not Path(checkpoint).is_file():
132 raise ValueError('GROVER requires an existing --checkpoint (DeepChem component format)')
133 if not isinstance(config, dict) or 'hidden_size' not in config:
134 raise ValueError('GROVER config must declare the checkpoint hidden_size')
135 allowed = {'hidden_size', 'num_attn_heads', 'depth', 'dropout', 'activation',
136 'self_attention', 'attn_out_size', 'ffn_num_layers', 'ffn_hidden_size'}
137 if set(config) - allowed:
138 raise ValueError(f'Unsupported GROVER config keys: {sorted(set(config) - allowed)}')
139 model = GroverModel(node_fdim=151, edge_fdim=165, features_dim=2048,
140 task='finetuning', mode=task_type, n_tasks=n_tasks,
141 n_classes=2 if task_type == 'classification' else None,
142 batch_size=16, learning_rate=1e-4, device='cpu', **config)
143 state = torch.load(checkpoint, map_location='cpu', weights_only=True)
144 if not isinstance(state, dict) or 'embedding' not in state:
145 raise ValueError('Checkpoint has no DeepChem embedding component')
146 model.components['embedding'].load_state_dict(state['embedding'], strict=True)
147 return model
150def train_grover(train_dataset, test_dataset, task_type='classification',
151 n_tasks=1, n_epochs=20, checkpoint=None, config=None,
152 valid_dataset=None):
153 valid_dataset = valid_dataset if valid_dataset is not None else test_dataset
154 validate_splits((train_dataset, valid_dataset, test_dataset), task_type, n_tasks)
155 model = build_grover(task_type, n_tasks, config, checkpoint)
156 model.fit(train_dataset, nb_epoch=n_epochs)
157 return model, evaluate_splits(model, (train_dataset, valid_dataset, test_dataset), task_type)
160def main():
161 parser = argparse.ArgumentParser(description=__doc__)
162 parser.add_argument('--model', choices=list(PRETRAINED_MODELS), required=True)
163 parser.add_argument('--dataset', choices=list(MOLNET_DATASETS))
164 parser.add_argument('--data')
165 parser.add_argument('--target', nargs='+', default=['target'])
166 parser.add_argument('--smiles-col', default='smiles')
167 parser.add_argument('--task-type', choices=['classification', 'regression'], default='classification')
168 parser.add_argument('--epochs', type=int, default=10)
169 parser.add_argument('--model-id', help='HF repository or local save_pretrained directory')
170 parser.add_argument('--revision', help='HF revision; remote code requires a reviewed immutable commit SHA')
171 parser.add_argument('--trust-remote-code', action='store_true', help='Enable model repository code after review')
172 parser.add_argument('--local-files-only', action='store_true')
173 parser.add_argument('--checkpoint', help='Local DeepChem GROVER component checkpoint')
174 parser.add_argument('--grover-config', help='JSON architecture matching the checkpoint')
175 args = parser.parse_args()
176 if not args.dataset and not args.data:
177 print('Error: Must specify either --dataset or --data', file=sys.stderr)
178 return 1
179 if args.dataset and args.data:
180 print('Error: Cannot specify both --dataset and --data', file=sys.stderr)
181 return 1
182 if args.epochs < 1:
183 parser.error('--epochs must be positive')
184 try:
185 # Fail before benchmark/weight downloads if required inputs are absent.
186 if args.model == 'grover' and (not args.checkpoint or not args.grover_config):
187 raise ValueError('GROVER requires --checkpoint and --grover-config')
188 if args.model == 'molformer' and (not args.trust_remote_code or not args.revision):
189 raise ValueError('MoLFormer requires reviewed --trust-remote-code and --revision')
190 if args.model != 'grover':
191 validate_remote_revision(args.trust_remote_code, args.revision)
192 if args.dataset:
193 task_type, n_tasks = MOLNET_DATASETS[args.dataset]
194 if args.model != 'grover' and n_tasks != 1:
195 raise ValueError('HF example supports one task; choose a single-task dataset')
196 tasks, datasets, _ = load_molnet_dataset(args.dataset, args.model)
197 train, valid, test = datasets
198 n_tasks = len(tasks)
199 else:
200 train, valid, test = load_custom_dataset(args.data, args.target, args.smiles_col, args.model)
201 task_type, n_tasks = args.task_type, len(args.target)
202 if args.model == 'grover':
203 config = json.loads(Path(args.grover_config).read_text())
204 train_grover(train, test, task_type, n_tasks, args.epochs,
205 args.checkpoint, config, valid)
206 else:
207 train_hf(args.model, train, valid, test, task_type, n_tasks, args.epochs,
208 model_id=args.model_id, revision=args.revision,
209 trust_remote_code=args.trust_remote_code,
210 local_files_only=args.local_files_only)
211 print('[OK] Fine-tuning finished; evaluate applicability before using predictions')
212 return 0
213 except Exception as error:
214 print(f'[FAIL] {error}', file=sys.stderr)
215 return 1
218if __name__ == '__main__':
219 sys.exit(main())