Help fix this issue
import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm


from transformers import T5ForConditionalGeneration, T5Tokenizer

tokenizer = T5Tokenizer.from_pretrained('t5-small')
model = T5ForConditionalGeneration.from_pretrained('t5-small')

optimizer = optim.AdamW(model.parameters(), lr=2e-5)

criterion = nn.CrossEntropyLoss(ignore_index=-100)

num_epochs = 1
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

best_val_loss = float('inf')
best_epoch = -1

for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0.0
    progress_bar = tqdm(enumerate(train_loader), total=len(train_loader))

    for batch_idx, batch in progress_bar:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()

        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()

        optimizer.step()

        epoch_loss += loss.item()
        progress_bar.set_description(f'Epoch {epoch + 1}/{num_epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}')

    average_train_loss = epoch_loss / len(train_loader)
    print(f'Epoch {epoch + 1}/{num_epochs}, Average Training Loss: {average_train_loss:.4f}')

    # 
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss

            val_loss += loss.item()

    average_val_loss = val_loss / len(val_loader)
    print(f'Epoch {epoch + 1}/{num_epochs}, Average Validation Loss: {average_val_loss:.4f}')

    if average_val_loss < best_val_loss:
        best_val_loss = average_val_loss
        best_epoch = epoch + 1
        torch.save(model.state_dict(), f'question_generation_model_epoch_{epoch + 1}.pth')
        print(f'Saved checkpoint at epoch {epoch + 1}')

print(f'Best validation loss: {best_val_loss:.4f} at epoch {best_epoch}')
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
Epoch 1/1, Batch 26999/27000, Loss: 1.9276: 100%|██████████| 27000/27000 [1:05:11<00:00,  6.90it/s]Epoch 1/1, Average Training Loss: 1.3072

Epoch 1/1, Average Validation Loss: 1.1229
Saved checkpoint at epoch 1
Best validation loss: 1.1229 at epoch 1
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer


model = T5ForConditionalGeneration.from_pretrained('t5-small')


model.load_state_dict(torch.load('/content/question_generation_model_epoch_1.pth', map_location=torch.device('cpu')))
model.eval()

def generate_question_or_answer(context, task):
    # Determine the task and format the input text accordingly
    if task.lower() == 'question':
        input_text = f"[QUESTION] [SEP] {context}"
    elif task.lower() == 'answer':
        input_text = f"[ANSWER] [SEP] {context}"
    else:
        raise ValueError("Task must be 'question' or 'answer'.")

    # Tokenize and encode the input text
    inputs = tokenizer.encode(input_text, return_tensors='pt', max_length=512, truncation=True)
    
    # Generate output
    with torch.no_grad():
        outputs = model.generate(inputs.to(device),
                                 max_length=128,  # Adjust max_length based on expected output length
                                 num_beams=4,     # Adjust num_beams for beam search
                                 early_stopping=True)  # Enable early stopping to stop when best beam reaches EOS
    
    # Decode the generated output and skip special tokens
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return generated_text


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

user_input = """
Tôi là con bò sống ở đồng quê, tôi yêu nhà cửa và công chúa
"""

# Generate question
question = generate_question_or_answer(user_input, 'question')

# Generate answer
answer = generate_question_or_answer(user_input, 'answer')

print(f"Generated Question: {question}")
print(f"Generated Answer: {answer}")
print()
Generated Question: [SEP] Tôi yêu nhà ca và công cha là g?
Generated Answer: [SEP] Tôi yêu nhà ca và công cha là g?