关于 HF 的 transformers — position_ids 和 padding_side 的注解。 我们可以通过 1) 左填充和 2) 传入适当的 position_ids ,使用 Hugging Face 的 transformer 复制精确的 logits:
import torch
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2", padding_side="right")
tokenizer.add_special_tokens({"pad_token": "[PAD]"})
pad_id = tokenizer.pad_token_id
query = torch.tensor([
[pad_id, pad_id, 23073],
])
response = torch.tensor([
[11, 339, 561],
])
temperature = 1.0
query = torch.tensor(query)
response = torch.tensor(response).long()
context_length = query.shape[1]
query_response = torch.cat((query, response), 1)
pretrained_model = transformers.AutoModelForCausalLM.from_pretrained("gpt2")
def forward(policy, query_responses, tokenizer):
attention_mask = query_responses != tokenizer.pad_token_id
position_ids = attention_mask.cumsum(1) - attention_mask.long() # exclusive cumsum
input_ids = query_responses.clone()
input_ids[~attention_mask] = 0
return policy(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
return_dict=True,
output_hidden_states=True,
)
output = forward(pretrained_model, query_response, tokenizer)
logits = output.logits
logits /= temperature
print(logits)
"""
tensor([[[ -26.9395, -26.4709, -30.0456, ..., -33.2208, -33.2884,
-27.4360],
[ -27.1677, -26.7330, -30.2386, ..., -33.6813, -33.6931,
-27.5928],
[ -35.2869, -34.2875, -38.1608, ..., -41.5958, -41.0821,
-35.3658],
[-111.3040, -110.9447, -112.9062, ..., -113.1306, -113.7788,
-109.1734],
[-111.5152, -109.6108, -114.9024, ..., -118.4352, -111.5668,
-112.1248],
[-122.6978, -121.8447, -128.2742, ..., -132.2805, -130.3961,
-125.7078]]], grad_fn=<DivBackward0>)
"""