Entrene el modelo con el siguiente código: import torch
import torch.nn as nn
import time
from torchtext.datasets import PennTreebank
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. Definir el modelo
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_layers)
        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 = to_map_style_dataset(PennTreebank(split='train'))
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>'])

# Dividir los datos en conjuntos de entrenamiento y validación
train_ratio = 0.8
train_data_size = int(len(train_data_raw) * train_ratio)
train_data = train_data_raw[:train_data_size]
valid_data = train_data_raw[train_data_size:]


# 3. Entrenar el modelo
def generate_pairs(tokens, shift=1):
    source = [tokens[i] for i in range(0, len(tokens) - shift)]
    target = [tokens[i] for i in range(shift, len(tokens))]
    return source, target

def tokenize_data(text, max_seq_len=None):
    tokens = tokenizer(text)[:max_seq_len]
    return torch.tensor([vocab[token] if token in vocab else vocab['<unk>'] for token in tokens], dtype=torch.long)

def create_batches(data, batch_size, max_seq_len=None):
    batched_data = []
    for i in range(0, len(data), batch_size):
        batched_text = []
        for text in data[i:i + batch_size]:
            tokens = tokenize_data(text, max_seq_len=max_seq_len)
            if len(tokens) <= 1:
                continue
            batched_text.append(tokens)
        batched_data.append(batched_text)
    return batched_data

def collate_fn(batch):
    batch_srcs = []
    batch_tgts = []
    max_len = max([len(text) for text in batch])
    for text in batch:
        source, target = generate_pairs(text)
        source = torch.tensor(source, dtype=torch.long)
        target = torch.tensor(target, dtype=torch.long)
        source = torch.cat([source, torch.full((max_len - len(source),), vocab['<pad>'], dtype=torch.long)], dim=0)
        target = torch.cat([target, torch.full((max_len - len(target),), vocab['<pad>'], dtype=torch.long)], dim=0)
        batch_srcs.append(source)
        batch_tgts.append(target)
    return torch.stack(batch_srcs).T, torch.stack(batch_tgts).T


vocab_size = len(vocab)
d_model = 512
nhead = 8  # Ajuste de hparámetro
num_layers = 6  # Ajuste de hiperparámetro
num_epochs = 8
learning_rate = 5e-4  # Ajuste de hiperparámetro
batch_size = 32  # Ajuste de hiperparámetro
max_seq_len = 50  # Limitar la longitud máxima de la secuencia para evitar problemas de memoria
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_batches = create_batches(train_data, batch_size, max_seq_len=max_seq_len)
valid_batches = create_batches(valid_data, 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).to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

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}")
Este es el resultado: Epoch: 1,	Loss: 9.557628631591797
Epoch: 1,	Loss: 3.6078901290893555
Epoch: 1,	Average Loss: 3.797003845298245,	Time taken: 172.61485052108765 seconds
Epoch: 2,	Loss: 3.8685905933380127
Epoch: 2,	Loss: 3.5860283374786377
Epoch: 2,	Average Loss: 3.754856805837653,	Time taken: 176.50595712661743 seconds
Epoch: 3,	Loss: 3.851348876953125
Epoch: 3,	Loss: 3.585386276245117
Epoch: 3,	Average Loss: 3.7396161551711224,	Time taken: 178.4089593887329 seconds
Epoch: 4,	Loss: 3.8236441612243652
Epoch: 4,	Loss: 3.579303741455078
Epoch: 4,	Average Loss: 3.7357346311268245,	Time taken: 178.76842260360718 seconds
Epoch: 5,	Loss: 3.81239652633667
Epoch: 5,	Loss: 3.576066255569458
Epoch: 5,	Average Loss: 3.7307242331849304,	Time taken: 178.6585669517517 seconds
Epoch: 6,	Loss: 3.808555841445923
Epoch: 6,	Loss: 3.56101655960083
Epoch: 6,	Average Loss: 3.7228062785170377,	Time taken: 178.71051454544067 seconds
Epoch: 7,	Loss: 3.795790910720825
Epoch: 7,	Loss: 3.565946340560913
Epoch: 7,	Average Loss: 3.716420245034613,	Time taken: 178.78500962257385 seconds
Epoch: 8,	Loss: 3.7901322841644287
Epoch: 8,	Loss: 3.561314821243286
Epoch: 8,	Average Loss: 3.711263506358114,	Time taken: 178.79905247688293 seconds
Epoch: 1,	Loss: 3.7836077213287354
Epoch: 1,	Loss: 3.5541670322418213
Epoch: 1,	Average Train Loss: 3.7087759117209864,	Average Valid Loss: 3.8043905900005153,	Time taken: 192.19864773750305 seconds
Epoch: 2,	Loss: 3.7849230766296387
Epoch: 2,	Loss: 3.5555243492126465
Epoch: 2,	Average Train Loss: 3.706186919157949,	Average Valid Loss: 3.790437877858093,	Time taken: 192.1538918018341 seconds
Epoch: 3,	Loss: 3.785423994064331
Epoch: 3,	Loss: 3.55546236038208
Epoch: 3,	Average Train Loss: 3.7053607400832522,	Average Valid Loss: 3.7913252790617853,	Time taken: 192.08488297462463 seconds
Epoch: 4,	Loss: 3.7783796787261963
Epoch: 4,	Loss: 3.5582540035247803
Epoch: 4,	Average Train Loss: 3.70311323046231,	Average Valid Loss: 3.7900776029086387,	Time taken: 191.91651272773743 seconds
Epoch: 5,	Loss: 3.781745433807373
Epoch: 5,	Loss: 3.5588197708129883
Epoch: 5,	Average Train Loss: 3.704557405451858,	Average Valid Loss: 3.7895773001043516,	Time taken: 192.08939933776855 seconds
Epoch: 6,	Loss: 3.778266429901123
Epoch: 6,	Loss: 3.569171905517578
Epoch: 6,	Average Train Loss: 3.706234157538686,	Average Valid Loss: 3.7815214363794363,	Time taken: 192.02809739112854 seconds
Epoch: 7,	Loss: 3.780600070953369
Epoch: 7,	Loss: 3.568856716156006
Epoch: 7,	Average Train Loss: 3.7157011594155898,	Average Valid Loss: 3.7747941189392438,	Time taken: 191.90934348106384 seconds
Epoch: 8,	Loss: 3.7852725982666016
Epoch: 8,	Loss: 3.5627803802490234
Epoch: 8,	Average Train Loss: 3.7210256421067416,	Average Valid Loss: 3.773217993544082,	Time taken: 191.81168031692505 seconds
Predicted output: <pad> <pad>