fineneo / finetune.py
tekkonetes's picture
Create finetune.py
af85fa9
raw
history blame
890 Bytes
mport torch
from transformers import GPTNeoForCausalLM, GPT2Tokenizer
# Load the pre-trained GPT-Neo model tokenizer (Replace with smaller or larger version)
model = GPTNeoForCausalLM.from_pretrained('EleutherAI/gpt-neo-1.3B')
tokenizer = GPT2Tokenizer.from_pretrained('EleutherAI/gpt-neo-1.3B')
# Load the text dataset
with open('text.txt', 'r') as f:
text = f.read()
# Tokenize the text dataset
encoding = tokenizer.encode(text, return_tensors='pt')
# Fine-tune the model on the text dataset
model.train()
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
for i in range(100):
loss = model(encoding, labels=encoding)[0]
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f'Epoch {i+1}, Loss: {loss.item()}')
# Save the fine-tuned model and tokenizer
model.save_pretrained('fine-tuned-gpt-neo')
tokenizer.save_pretrained('fine-tuned-gpt-neo')