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 get_tokenizer(self): | |
return self.tokenizer | |