|
from transformers import GPT2LMHeadModel, GPT2Tokenizer |
|
import torch |
|
|
|
|
|
tokenizer = GPT2Tokenizer.from_pretrained('gpt2') |
|
model = GPT2LMHeadModel.from_pretrained('gpt2') |
|
|
|
|
|
with open('train.txt', 'r') as f: |
|
text = f.read() |
|
|
|
|
|
input_ids = tokenizer.encode(text, return_tensors='pt') |
|
|
|
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5) |
|
model.train() |
|
for i in range(100): |
|
outputs = model(input_ids, labels=input_ids) |
|
loss = outputs[0] |
|
loss.backward() |
|
optimizer.step() |
|
optimizer.zero_grad() |
|
|
|
print(f'Epoch {i+1}, Loss: {loss.item()}') |
|
|
|
|
|
model.save_pretrained('my_gpt_model') |
|
|