| from transformers import pipeline, AutoTokenizer | |
| class gpt2: | |
| def __init__(self,device="cpu"): | |
| self.text_generation = pipeline("text-generation", model="gpt2",device=device) | |
| self.tokenizer = AutoTokenizer.from_pretrained("gpt2") | |
| def generate_text(self,**kwargs): | |
| results = self.text_generation(**kwargs) | |
| return [item['generated_text'] for item in results[0]] | |
| def get_tokenizer(self): | |
| return self.tokenizer | |
| if __name__ == '__main__': | |
| gpt2 = gpt2() | |
| print(gpt2.generate_text(["Hello, how are you?","I am fine, thank you."])) | |