gpt-pretrain/generate.py

91 lines
2.7 KiB
Python

import argparse
from typing import List
import torch
from transformers import PreTrainedModel, PreTrainedTokenizer
from utils import load_model, load_tokenizer
def eval_prompts(
model: PreTrainedModel,
tokenizer: PreTrainedTokenizer,
prompts: List[str],
use_tril_attention_mask: bool = False,
) -> List[str]:
inputs = tokenizer(
prompts, padding=True, return_tensors='pt', return_attention_mask=True
)
inputs['position_ids'] = inputs.attention_mask.cumsum(-1) - 1
inputs['position_ids'].masked_fill_(inputs.attention_mask == 0, 1)
if use_tril_attention_mask:
inputs['attention_mask'] = (
inputs.attention_mask.unsqueeze(1) * inputs.attention_mask.unsqueeze(2)
).tril()
inputs = inputs.to(model.device)
with torch.inference_mode():
output_ids = model.generate(
**inputs,
do_sample=False,
num_beams=16,
max_new_tokens=100,
eos_token_id=tokenizer.eos_token_id,
early_stopping=True,
)
completes = tokenizer.batch_decode(
output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
return completes
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model_name_or_path",
type=str,
help="Name of or path to model",
default='gpt2',
)
parser.add_argument(
"--use_tril_attention_mask",
help="Use tril attention mask during training",
action="store_true",
)
parser.add_argument(
"--tokenizer_name_or_path",
type=str,
help="Name of or path to tokenizer",
default=None,
)
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
if args.tokenizer_name_or_path is None:
args.tokenizer_name_or_path = args.model_name_or_path
device = torch.device(0)
model = load_model(args.model_name_or_path)
tokenizer = load_tokenizer(args.tokenizer_name_or_path)
model = model.to(device)
prompts = [
"Shall I compare thee to a summer's day? Thou art more lovely and more temperate.",
"Shall I compare thee to a summer's day? Thou art more lovely and",
"Belle! C'est un mot qu'on dirait inventé pour elle.",
"Belle! C'est un mot qu'on dirait inventé",
"这是一个最好的时代,这是一个最坏的时代。",
"这是一个最好的时代,这是一个最坏的",
]
completes = eval_prompts(
model, tokenizer, prompts, use_tril_attention_mask=args.use_tril_attention_mask
)
for prompt, complete in zip(prompts, completes):
print("[p]", prompt)
print("[c]", complete)