from transformers import GTP2Tokenizer, TrainingArguments, Trainer, GPT2LMHeadModel
from datasets import load_dataset

dataset = load_dataset("sst2")

for row in dataset['train']:
    print(row)

for i, row in enumerate(dataset):
    prep_text = f"<|startoftext|> {rwo['sentence']}<|pad|>Sentiment: {rwo['label']}<|endoftext|>"
    encodings_dict = tokenizer(prep_txt)

tokenizer = GTP2Tokenizer.from_pretrained('gpt2', bos_token='<|startoftext|>', eos_token='<|endoftext|>', pad_token='<|pad|>')
model = GPT2LMHeadModel.from_pretrained('gpt2')

train_args = TrainingArguments(output_dir='results', num_train_epochs = 1, warmup_steps =100, weight_decay = 0.01)

Trainer(model='gpt2', args=train_args,train_dataset=train_dataset)

model.eval()

prompt = f'<|startoftext|>Tweet: {text}\nSentiment:'
tokenizer_text = tokenizer(prompt, return_tensors="pt").input_ids
output = model.generate(tokenized_text)
predicted_text = tokenizer.decode(output)