HuggingFace 참고
https://huggingface.co/ctheodoris/Geneformer/tree/main
ctheodoris/Geneformer at main
huggingface.co
요구사항은 requirements.txt 참고(패키지 설치)
모델 관련 폴더
Geneformer/
├── Geneformer-V1-10M
├── Geneformer-V2-104M
├── Geneformer-V2-104M_CLcancer // 암 세포 데이터에 대해 특화된 pretrained 버전으로 보임
├── Geneformer-V2-316M
pretrained 된 모델 가중치가 저장됨.
V1, V2는 버전, 10M, 104M..등은 파라미터 개수
<각 폴더 내부 구조>
config.json: 모델 구조 설정(hidden_size, hidden_dropout_prob, model_type, num_attention_heads 등)
model.safetensors: pytorch_model.bin의 안전한 버전
pytorch_model.bin: pytorch 형식의 실제 모델 weight 파일
training_args.bin: 모델 사전학습시에 사용된 하이퍼파라미터
모델 아키텍쳐, 코드
cf. .pkl 파일: python pickle 포맷, 직렬화된 데이터를 담는 바이너리 파일
geneformer/ (일부만 표시)
├── gene_dictionaries_30m // v1 모델 -> Geneformer V1 10M에 매핑?되는 파일들
├── ensembl_mapping_dict_gc30M.pkl // 유전자ID <-> 기타 ID 매핑 정보
├── gene_median_dictionary_gc30M.pkl // 각 유전자의 median expression 정보(정규화에 사용)
├── gene_name_id_dict_gc30M.pkl // 유전자 이름 <-> 내부 ID 매핑
├── token_dictionary_gc30M.pkl // 모델에 입력되는 토큰과 ID 간의 매핑
├── classifier.py // 모델 정의
├── classifier_utils.py // classifier 보조 함수
├── mtl_classifier.py // 멀티태스크 학습을 위한 classifier 모델 구조를 정의
├── tokenizer.py // rank encoding
├── emb_extractor.py // 임베딩 추출
├── collator_for_classification.py // 전처리 도우미(?)
├── perturber_utils.py // perturbation 실험용
├── in_silico_perturber.py // 유전자 삭제 실험
├── in_silico_perturber_stats.py // in_silico_perturber.py가 유전자를 삭제하거나 활성화한 뒤, 모델 임베딩의 변화량을 저장하면 이 파일이 그 차이를 정량화하고 통계적으로 유의한지 분석
├── pretrainer.py // 사전 학습
├── evaluation_utils.py // 성능 평가 함수
├── ensembl_mapping_dict_gc104M.pkl
├── gene_median_dictionary_gc104M.pkl // rank normalization 참조용
├── gene_name_id_dict_gc30M.pkl // 유전자 이름 <-> 내부 ID 매핑
├── token_dictionary_gc30M.pkl // 모델에 입력되는 토큰과 ID 간의 매핑
데이터 전처리 & rank encoding - tokenizer.py
gene vector 배열을 만들고, 발현량이 0이 아닌 유전자만 토큰화
rank_genes()에서 발현량이 높은 유전자가 시퀀스의 앞쪽에 위치하도록 토큰 정렬
def rank_genes(gene_vector, gene_tokens):
# sort by median-scaled gene values
sorted_indices = np.argsort(-gene_vector)
return gene_tokens[sorted_indices]
def tokenize_cell(gene_vector, gene_tokens):
# create array of gene vector with token indices
# mask undetected genes
nonzero_mask = np.nonzero(gene_vector)[0]
# rank by median-scaled gene values
return rank_genes(gene_vector[nonzero_mask], gene_tokens[nonzero_mask])
TranscriptomeTokenizer 클래스에서는, 파일형식이 .h5ad, .loom이냐에 따라 나누어서 토큰화 진행
실제 유전자 데이터 파일을 열고, 발현값 정규화 → 유전자별 rank 변환 → token list 생성
✨ 논문에서 ranking매길 시에 유전자들의 발현량을 절대적으로 비교하는게 아니라, 특정 유전자의 발현량을 전체 세포에서의 그 유전자의 발현량의 정규화값(평균?중앙값?)과 비교해서 상대적으로 ranking한다고 언급 ... → 코드에서 확인 완료
class TranscriptomeTokenizer:
# 생략
def tokenize_anndata(self, adata_file_path, target_sum=10_000):
adata = sum_ensembl_ids(
adata_file_path,
self.collapse_gene_ids,
self.gene_mapping_dict,
self.gene_token_dict,
self.custom_attr_name_dict,
self.use_h5ad_index,
file_format="h5ad",
chunk_size=self.chunk_size,
)
# 생략
for i in range(0, len(filter_pass_loc), self.chunk_size):
idx = filter_pass_loc[i : i + self.chunk_size]
n_counts = adata[idx].obs["n_counts"].values[:, None]
X_view0 = adata[idx, :].X
X_view = X_view0[:, coding_miRNA_loc]
X_norm_unscaled = X_view / n_counts * target_sum
# X_norm_unscaled: 각 세포의 유전자 발현값을 정규화
# norm_factor_vector: gene_median_dict에서 가져온 각 유전자의 전체 cell에서의 발현량 median
X_norm = X_norm_unscaled / norm_factor_vector
X_norm = sp.csr_matrix(X_norm)
X_norm_unscaled = sp.csr_matrix(X_norm_unscaled)
tokenized_cells += [
rank_genes(X_norm[i].data, coding_miRNA_tokens[X_norm[i].indices])
for i in range(X_norm.shape[0])
]
if self.keep_counts:
X_norm_unscaled = sp.csr_matrix(X_norm_unscaled)
tokenized_counts += [
rank_genes(X_norm[i].data, X_norm_unscaled[i].data)
for i in range(X_norm.shape[0])
]
# add custom attributes for subview to dict
if self.custom_attr_name_dict is not None:
for k in file_cell_metadata.keys():
file_cell_metadata[k] += adata[idx].obs[k].tolist()
else:
file_cell_metadata = None
모델 아키텍쳐 - classifier.py
classifier는 cell, gene으로 나뉨.
사용 예시는 아래 Examples 참고
아래는 model을 불러오고 training 파라미터 불러오는 예시. + fine-tuning
from . import perturber_utils as pu
##### Load model and training args #####
model = pu.load_model(
self.model_type,
num_classes,
model_directory,
"train",
quantize=self.quantize,
)
def_training_args, def_freeze_layers = cu.get_default_train_args(
model, self.classifier, train_data, output_directory
)
del model
if self.training_args is not None:
def_training_args.update(self.training_args)
logging_steps = round(
len(train_data) / def_training_args["per_device_train_batch_size"] / 10
)
def_training_args["logging_steps"] = logging_steps
def_training_args["output_dir"] = output_directory
if eval_data is None:
def_training_args["evaluation_strategy"] = "no"
def_training_args["load_best_model_at_end"] = False
if transformers_version >= parse("4.46"):
def_training_args["eval_strategy"] = def_training_args.pop("evaluation_strategy")
def_training_args.update(
{"save_strategy": "epoch", "save_total_limit": 1}
) # only save last model for each run
training_args_init = TrainingArguments(**def_training_args)
##### Fine-tune the model #####
# define the data collator
if self.classifier == "cell":
data_collator = DataCollatorForCellClassification(
token_dictionary=self.gene_token_dict
)
elif self.classifier == "gene":
data_collator = DataCollatorForGeneClassification(
token_dictionary=self.gene_token_dict
)
pertuber_utils.py
from transformers import (
BertForMaskedLM,
BertForSequenceClassification,
BertForTokenClassification,
BitsAndBytesConfig,
)
# load model to GPU
def load_model(model_type, num_classes, model_directory, mode, quantize=False):
if model_type == "Pretrained-Quantized":
inference_only = True
model_type = "Pretrained"
quantize = True
elif model_type == "MTLCellClassifier-Quantized":
inference_only = True
model_type = "MTLCellClassifier"
quantize = True
else:
inference_only = False
output_hidden_states = (mode == "eval")
# Quantization logic
if isinstance(quantize, dict):
quantize_config = quantize.get("bnb_config", None)
peft_config = quantize.get("peft_config", None)
elif quantize:
if inference_only:
quantize_config = BitsAndBytesConfig(load_in_8bit=True)
peft_config = None
else:
quantize_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
try:
peft_config = LoraConfig(
lora_alpha=128,
lora_dropout=0.1,
r=64,
bias="none",
task_type="TokenClassification",
)
except ValueError as e:
peft_config = LoraConfig(
lora_alpha=128,
lora_dropout=0.1,
r=64,
bias="none",
task_type="TOKEN_CLS",
)
else:
quantize_config = None
peft_config = None
# Model class selection
model_classes = {
"Pretrained": BertForMaskedLM,
"GeneClassifier": BertForTokenClassification,
"CellClassifier": BertForSequenceClassification,
"MTLCellClassifier": BertForMaskedLM
}
model_class = model_classes.get(model_type)
if not model_class:
raise ValueError(f"Unknown model type: {model_type}")
# Model loading
model_args = {
"pretrained_model_name_or_path": model_directory,
"output_hidden_states": output_hidden_states,
"output_attentions": False,
}
if model_type != "Pretrained":
model_args["num_labels"] = num_classes
if quantize_config:
model_args["quantization_config"] = quantize_config
# Load the model
model = model_class.from_pretrained(**model_args)
if mode == "eval":
model.eval()Transformer 구조 기반인 BERT를 로딩하고 있는 것으로 확인됨
BertModel ← BERT base (encoder-only transformer)
├── BertForMaskedLM ← [BERT + MLM head] → Pretraining용
├── BertForSequenceClassification ← [BERT + classification head]
├── BertForTokenClassification ← [BERT + token classification head]
⁉️BERT에 대해서도 공부해봐야할듯...
Examples(downstream task 적용)
examples/
├── cell_classification.ipynb // 세포단위로 질병 상태 분류
├── distributed_multitask_cell_classification.ipynb // multi-task learning ex) 세포 타입, 질병 상태 동시에 분류
├── extract_and_plot_cell_embeddings.ipynb // 임베딩 추출, 시각화 예제 코드
├── gene_classification.ipynb // dosage sensitive vs insensitive TFs 분류(논문 참고)
├── in_silico_pertubation_ipynb // 어떤 유전자의 삭제가 dcm -> 정상 상태로 임베딩을 shifting 시키는 지 알아내는 예제
├── multitask_cell_classification.ipynb // how to use the Geneformer multi-task cell classifier
├── tokenizing_scRNAseq_data.ipynb // .loom, .h5ad scRNA-seq 데이터를 rank value encoding .dataset으로 토큰화하는 코드
<cell_classification.ipynb>
cardiomyopathy 질병 상태(nf/hcm/dcm) 를 세포 단위에서 예측하도록 Geneformer를 fine-tune하는 예시
from geneformer import Classifier
cc = Classifier(classifier="cell",
cell_state_dict = {"state_key": "disease", "states": "all"},
filter_data=filter_data_dict,
training_args=training_args,
max_ncells=None,
freeze_layers = 2,
num_crossval_splits = 1,
forward_batch_size=200,
model_version="V1", # OF NOTE: SET TO V1 MODEL, PROVIDE V1 MODEL PATH IN SUBSEQUENT CODE
nproc=16)먼저, geneformer에서 Classifier 불러오고,
Classifier 불러올 때 classifier="cell"로 정의
이때, 학습 파라미터는 pretrained된 값 사용
cc.prepare_data(input_data_file="/path/to/human_dcm_hcm_nf_2048_w_length.dataset",
output_directory=output_dir,
output_prefix=output_prefix,
split_id_dict=train_test_id_split_dict)데이터 준비: dataset을 train/valid/test로 나누고 전처리 진행
all_metrics = cc.validate(model_directory="/path/to/Geneformer", # OF NOTE: SET TO V1 MODEL ABOVE, PROVIDE V1 MODEL PATH HERE
prepared_input_data_file=f"{output_dir}/{output_prefix}_labeled_train.dataset",
id_class_dict_file=f"{output_dir}/{output_prefix}_id_class_dict.pkl",
output_directory=output_dir,
output_prefix=output_prefix,
split_id_dict=train_valid_id_split_dict)사전학습모델(v1 혹은 v2) 사용하여 학습 + 검증 진행
즉, fine-tuning 진행
all_metrics_test = cc.evaluate_saved_model(
model_directory=f"{output_dir}/{datestamp_min}_geneformer_cellClassifier_{output_prefix}/ksplit1/",
id_class_dict_file=f"{output_dir}/{output_prefix}_id_class_dict.pkl",
test_data_file=f"{output_dir}/{output_prefix}_labeled_test.dataset",
output_directory=output_dir,
output_prefix=output_prefix,
)테스트 성능 평가 (evaluate_saved_model())로 테스트셋 평가)
<gene_classification.ipynb>
Dosage-sensitive vs. -Insensitive Transcription Factors(TFs)
from geneformer import Classifier
cc = Classifier(classifier="gene",
gene_class_dict = gene_class_dict,
max_ncells = 10_000,
freeze_layers = 4,
num_crossval_splits = 5,
forward_batch_size=200,
model_version="V1", # OF NOTE: SET TO V1 MODEL, PROVIDE V1 MODEL PATH IN SUBSEQUENT CODE
nproc=16)먼저, geneformer에서 Classifier 불러오고,
Classifier 불러올 때 classifier="gene"로 정의
이때, 학습 파라미터는 pretrained된 값 사용
cc.prepare_data(input_data_file="/path/to/gc-30M_sample50k.dataset",
output_directory=output_dir,
output_prefix=output_prefix)데이터 준비: dataset을 train/valid/test로 나누고 전처리 진행
# V1 model: https://huggingface.co/ctheodoris/Geneformer/blob/main/Geneformer-V1-10M/model.safetensors
all_metrics = cc.validate(model_directory="/path/to/Geneformer", # OF NOTE: SET TO V1 MODEL ABOVE, PROVIDE V1 MODEL PATH HERE
prepared_input_data_file=f"{output_dir}/{output_prefix}_labeled.dataset",
id_class_dict_file=f"{output_dir}/{output_prefix}_id_class_dict.pkl",
output_directory=output_dir,
output_prefix=output_prefix)사전학습모델(v1 혹은 v2) 사용하여 학습 + 검증 진행
즉, fine-tuning 진행; 이 예제에서는 5-fold cross-validation 진행
