File size: 310 Bytes
44466c7
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
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