基于MindFormers实现GPT2模型的推理

拾回程序猿的圈圈∞2024-08-07 23:39

前言

针对MindFormers的安装，可参考本专栏里的另一篇博客

pipeline方式

python 复制代码

from mindformers import pipeline
from mindformers import GPT2LMHeadModel, GPT2Config, GPT2Tokenizer
tok = GPT2Tokenizer.from_pretrained("gpt2")
gpt_model = GPT2LMHeadModel.from_pretrained("gpt2")
prompt="I love Shanghai"
text_generation = pipeline(task='text_generation', model=gpt_model, tokenizer=tok, max_length=50)
outputs=text_generation(prompt, top_k=3)
print(outputs)

持续更新中...