USER
import torch
import torch.nn as nn
from torchtext.datasets import WikiText2
from torchtext.data.functional import to_map_style_dataset
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data import DataLoader
# 1. Define the model
class TransformerModel(nn.Module):
def init(self, vocab_size, d_model, nhead, num_layers, dropout=0.1):
super(TransformerModel, self).init()
self.embedding = nn.Embedding(vocab_size, d_model)
self.transformer = nn.Transformer(d_model, nhead, num_layers, dropout=dropout)
self.fc = nn.Linear(d_model, vocab_size)
def forward(self, src, tgt):
src = self.embedding(src)
tgt = self.embedding(tgt)
x = self.transformer(src, tgt)
x = self.fc(x)
return x
# 2. Preparar datos
train_data_raw, valid_data_raw, _ = WikiText2(split=('train', 'valid', 'test'))
train_data_raw = to_map_style_dataset(train_data_raw)
valid_data_raw = to_map_style_dataset(valid_data_raw)
tokenizer = get_tokenizer("spacy", "en_core_web_sm")
# Agregar tokens especiales (UNK y PAD) al vocabulario
specials = ['<unk>','<pad>']
vocab = build_vocab_from_iterator((tokenizer(y) for y in train_data_raw), specials=specials)
vocab.set_default_index(vocab['<unk>'])
vocab_size = len(vocab)
d_model = 512 # Reduced to 512
nhead = 8
num_layers = 4 # Reduced to 4
num_epochs = 50 # Increased to 50
learning_rate = 5e-4 # Reduced to 5e-4
optimizer_choice = 'Adam'
batch_size = 128 # Increased to 128
max_seq_len = 100 # Increased to 100
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_batches = create_batches(train_data_raw, batch_size, max_seq_len=max_seq_len)
valid_batches = create_batches(valid_data_raw, batch_size, max_seq_len=max_seq_len)
train_loader = DataLoader(train_batches, batch_size=None, collate_fn=collate_fn)
valid_loader = DataLoader(valid_batches, batch_size=None, collate_fn=collate_fn)
model = TransformerModel(vocab_size=vocab_size, d_model=d_model, nhead=nhead, num_layers=num_layers, dropout=0.1).to(device)
loss_function = nn.CrossEntropyLoss()
# Use different optimizers
if optimizer_choice == "Adam":
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
elif optimizer_choice == "RMSprop":
optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate)
elif optimizer_choice == "Adagrad":
optimizer = torch.optim.Adagrad(model.parameters(), lr=learning_rate)
else:
raise ValueError("Invalid optimizer choice")
def train(model, epoch):
model.train()
train_loss = 0
for i, (source, target) in enumerate(train_loader):
source = source.to(device)
target = target.to(device)
optimizer.zero_grad()
output = model(source, target[:-1])
loss = loss_function(output.reshape(-1, vocab_size), target[1:].reshape(-1))
loss.backward()
optimizer.step()
train_loss += loss.item()
if i % 1000 == 0:
print(f'Epoch: {epoch + 1},\tLoss: {loss.item()}')
return train_loss / len(train_loader)
for epoch in range(num_epochs):
start_time = time.time()
train_loss = train(model, epoch)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Epoch: {epoch + 1},\tAverage Loss: {train_loss},\tTime taken: {elapsed_time} seconds")
# 4. Evaluar el modelo en el conjunto de validación
def evaluate(model, valid_loader):
model.eval()
valid_loss = 0
with torch.no_grad():
for source, target in valid_loader:
source = source.to(device)
target = target.to(device)
output = model(source, target[:-1])
loss = loss_function(output.reshape(-1, vocab_size), target[1:].reshape(-1))
valid_loss += loss.item()
return valid_loss / len(valid_loader)
# 5. Entrenar y evaluar el modelo en cada época
for epoch in range(num_epochs):
start_time = time.time()
train_loss = train(model, epoch)
valid_loss = evaluate(model, valid_loader)
end_time = time.time()
elapsed_time = end_time - start_time
print(f"Epoch: {epoch + 1},\tAverage Train Loss: {train_loss},\tAverage Valid Loss: {valid_loss},\tTime taken: {elapsed_time} seconds")
# 6. Guardar el modelo entrenado
torch.save(model.state_dict(), "trained_model.pth")
# 7. Cargar el modelo
def load_model(path, device):
loaded_model = TransformerModel(vocab_size=vocab_size, d_model=d_model, nhead=nhead, num_layers=num_layers).to(device)
loaded_model.load_state_dict(torch.load(path))
loaded_model.eval()
return loaded_model
loaded_model = load_model("trained_model.pth", device)
# 8. Predecir texto
def decode_output(output):
probabilities = torch.softmax(output, dim=2)
predicted_indices = torch.argmax(probabilities, dim=2)
return " ".join(vocab.lookup_tokens(predicted_indices.cpu().numpy().squeeze().tolist()))
def predict_next_token(input_text, model):
tokens = tokenize_data(input_text)
source, target = generate_pairs(tokens, shift=1)
source = torch.tensor(source, dtype=torch.long).unsqueeze(1).to(device)
target = torch.tensor(target, dtype=torch.long).unsqueeze(1).to(device)
with torch.no_grad():
output = model(source, target[:-1])
return decode_output(output)
input_text = "My favorite color is"
predicted_output = predict_next_token(input_text, loaded_model)
print(f"Predicted output: {predicted_output}")ASSISTANT
It looks like there is a mistake in the provided code. The `create_batches()` and `collate_fn()` functions are missing, which are essential for creating data loaders for training and validation.
You can implement the missing functions in the following way:
```python
from torchtext.data.functional import numericalize_tokens_from_iterator
def tokenize_data(data_raw):
return [torch.tensor(vocab(tokenizer(y))) for y in data_raw]
def generate_pairs(tokenized_data, shift):
source = tokenized_data[:-shift]
target = tokenized_data[shift:]
return source, target
def create_batches(data_raw, batch_size, max_seq_len=100):
tokenized_data = tokenize_data(data_raw)
pairs = [generate_pairs(tokenized_data[i:i+max_seq_len], shift=1) for i in range(0, len(tokenized_data)-max_seq_len, max_seq_len)]
batches = [pairs[i:i+batch_size] for i in range(0, len(pairs), batch_size)]
return batches
def collate_fn(batch):
sources, targets = zip(*batch)
src_max_len = max([len(x) for x in sources])
tgt_max_len = max([len(x) for x in targets])
padded_sources = torch.zeros(len(sources), src_max_len, dtype=torch.long)
padded_targets = torch.zeros(len(targets), tgt_max_len, dtype=torch.long)
for i, (source, target) in enumerate(zip(sources, targets)):
padded_sources[i, :len(source)] = source
padded_targets[i, :len(target)] = target
return padded_sources.t().contiguous(), padded_targets.t().contiguous()
```
Now your code should work fine when executing the training loop.
Keep in mind that the model will be very large and take a long time to train with the provided parameters. It would be recommended to reduce the number of layers, increase the batch size, or use a pre-trained model if available.