File size: 2,413 Bytes
c3d1a15 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import torch
from transformers import pipeline
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from google.colab import files
# device = torch.device("cpu" if torch.cuda.is_available() else "cuda")
device = torch.device("cpu")
# Load the GPT-2 model and tokenizer
model = GPT2LMHeadModel.from_pretrained("gpt2")
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
# Add a padding token to the tokenizer
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
# Move the model to the appropriate device
model = model.to(device)
# Upload your documents
uploaded_files = files.upload()
training_data = []
for file_name, content in uploaded_files.items():
try:
document = content.decode("utf-8")
except UnicodeDecodeError:
document = content.decode("latin-1")
training_data.append(document)
# Fine-tuning the GPT-2 model
tokenized_data = tokenizer('\n\n'.join(training_data), truncation=True, padding=True, max_length=256, return_tensors="pt")
model.resize_token_embeddings(len(tokenizer))
# Define the loss function
loss_function = torch.nn.CrossEntropyLoss()
# Define the training loop
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-5)
# optimizer = torch.nn.Module(optimizer)
# optimizer = optimizer.to(device)
model.train()
accumulation_steps = 4
batch_size = 4
for epoch in range(3):
for i in range(0, len(tokenized_data['input_ids']), batch_size):
input_ids_batch = tokenized_data['input_ids'][i:i + batch_size].to(device, non_blocking=True)
outputs = model(input_ids_batch)
logits = outputs.logits
loss = loss_function(logits.view(-1, model.config.vocab_size), input_ids_batch.view(-1))
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f"Epoch: {epoch+1}, Batch: {i+1}/{len(tokenized_data['input_ids'])//batch_size}, Loss: {loss.item()}")
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
# Start chatting with your trained model
while True:
user_input = input("User: ")
input_ids = tokenizer.encode(user_input, return_tensors='pt').to(device)
generate = pipeline('text-generation', model='gpt2')
output = model.generate(input_ids=input_ids, max_length=100, num_return_sequences=1)
response = tokenizer.decode(output[0])
print("ChatBot:", response)
break # Add a condition to break the while loop |