单细胞预训练模型工作总结(2)

gene_id = [gene_0, gene_1, gene_2, ..., gene_n]
exp_token = [count_0, count_1, count_2, ..., count_n] 

cell_embed = (gene_0_embed + count_0_embed, ..., gene_n_embed + count_n_embed) 

gene_id = [gene_0, gene_1, gene_2, ..., gene_n]
exp_token = [count_0, mask, count_2, ..., count_n] 

cell_embed = (gene_0_embed + count_0_embed, gene_1_embed + mask_embed..., gene_n_embed + count_n_embed) 

# GPT2 分词示例

# load the gpt-2 tokenizer
tokenizer = AutoTokenizer.from_pretrained("../hf_model/gpt2")
tokenizer.pad_token=tokenizer.eos_token 

# tokenize
def tokenize_function(example):
    #return tokenizer(text=example["text"]) # warning for seq-length larger than tokenizer.model_max_length ## 这里可以不需要截断的，后面的chunk函数会将文本拼接后分割成1024的输入长度
    return tokenizer(text=example['text'], truncation=True, max_length=tokenizer.model_max_length)
tokenized_ds = dataset.map(tokenize_function,batched=True,remove_columns='text', num_proc=num_proc)

# save to disk if required (use load_from_disk latter)
tokenized_ds.save_to_disk('../hf_data/gpt2tokenized_bert_pretrain')

# Make samples to a size of 1024
def concat(examples):    
    examples["input_ids"]=[list(chain.from_iterable(examples['input_ids']))] # convert chain to list of tokens
    examples["attention_mask"]=[list(chain.from_iterable(examples['attention_mask']))] # convert chain to list of tokens
    return examples
    
# takes a lot of time (worth saving it to disk)
concated_ds = tokenized_ds.map(concat,batched=True,batch_size=1000000,num_proc=256)

def chunk(examples):
    chunk_size = 1024 # modify this accordingly       
    input_ids = examples["input_ids"][0] # List[List], pass the inner list      
    attention_mask = examples["attention_mask"][0] # List[List]
    input_ids_truncated = []
    attention_mask_truncated = []
    
    #slice with step_size=chunk_size
    for i in range(0,len(input_ids),chunk_size):
        chunk = input_ids[i:i+chunk_size]
        if len(chunk)==chunk_size: # drop the last chunk if not equal
            input_ids_truncated.append(chunk)
            attention_mask_truncated.append(attention_mask[i:i+chunk_size])     
    examples['input_ids']=input_ids_truncated
    examples["attention_mask"]=attention_mask_truncated
        
    return examples   

chunked_ds = concated_ds.map(chunk,batched=True,batch_size=2,num_proc=256)
chunked_ds.save_to_disk('../hf_data/gpt2tokenized_bert_pretrain_chunked_ds') # will use this latter for diff experimentation

data_collator = DataCollatorForLanguageModeling(tokenizer,mlm=False)

# load the model
configuration = GPT2Config()
model =GPT2LMHeadModel(configuration)

# training arguments
training_args = TrainingArguments( output_dir='gpt-2-warm-up/standard-gpt',
                                  eval_strategy="steps",
                                  eval_steps=10,                                  
                                  num_train_epochs=1,
                                  per_device_train_batch_size=32,
                                  gradient_accumulation_steps=16,
                                  per_device_eval_batch_size=16,
                                  learning_rate=2.5e-4,
                                  lr_scheduler_type='cosine',
                                  warmup_ratio=0.05,
                                  adam_beta1=0.9,
                                  adam_beta2=0.999,                                  
                                  weight_decay=0.01,
                                  logging_dir='./logs',                                  
                                  logging_strategy="steps",
                                  logging_steps = 10,
                                  save_steps=1000,
                                  save_total_limit=10,                                  
                                  report_to='tensorboard'#'wandb',                                  
                                 ) 

trainer = Trainer(model=model,
                 args = training_args,
                 tokenizer=tokenizer,
                 train_dataset=chunked_ds["train"],
                 eval_dataset=chunked_ds["test"],
                 data_collator = data_collator)

prompts = 'A long time ago, there is a boy called Super Rui' #"I was telling her that"
inputs = tokenizer(prompts,return_tensors='pt').input_ids
outputs = model.generate(inputs, max_new_tokens=100, do_sample=True, top_k=10, top_p=0.95)
res = tokenizer.batch_decode(outputs, skip_special_tokens=True)
print(res)

A long time ago, there is a boy called Super Rui. He was a good man to be killed in a coma.
In the summer of 1925, the young man was taken prisoner to a prison and found a few days earlier. 
After the murder, the young man was shot to death in a suicide trial, a suicide sentence was 
found guilty of the murder, and a few were killed. The murder was released in a prison prison
in June 1921, and his wife died.The murder, the police was killed and the police arrived to

sentence = "Paris is the capital of <mask>."
inputs = tokenizer(sentence, return_tensors='pt')

outputs = model(**inputs)

# 获取预测的概率分布
predictions = outputs.logits

# 找到 mask 位置
mask_index = (inputs.input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]

# 获取 mask 位置的预测分布
predicted_distribution = predictions[0, mask_index].softmax(dim=1)  # 使用 softmax 转换为概率

# 获取概率最高的前十个 token ID 及其概率
top_k_values, top_k_indices = predicted_distribution.topk(10)

# 解码这些 token 并打印结果
for value, index in zip(top_k_values, top_k_indices):
    K = len(index)
    for i in range(K):
        idx = int(index[i].item())
        word = tokenizer.decode([idx])  # 将索引转换为单词
        probability = value[i].item()
        print(f"Word: {word}, Probability: {probability:.4f}")


Word:  me, Probability: 0.0416
Word:  you, Probability: 0.0397
Word:  hell, Probability: 0.0251
Word:  him, Probability: 0.0228
Word:  us, Probability: 0.0228
Word:  earth, Probability: 0.0191
Word:  town, Probability: 0.0185
Word:  her, Probability: 0.0161
Word:  them, Probability: 0.0153
Word:  yours, Probability: 0.015