Add a loop to the last section of the code where an input is passed to the model to generate text: import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset
from collections import Counter



class TextDataset(Dataset):
    def __init__(self, path, seq_len):
        self.seq_len = seq_len
        with open(path, "r", encoding="utf-8") as f:
            text = f.read()
        words = text.split()
        self.vocab, self.idx2token = self.build_vocab(words)
        self.tokens = [self.vocab.get(w, 0) for w in words]  # Default to 0 for unknown tokens

    def build_vocab(self, words):
        counts = Counter(words)
        vocab = {word: i for i, (word, _) in enumerate(counts.most_common(), 1)}
        vocab["<pad>"] = 0
        idx2token = {i: word for word, i in vocab.items()}
        return vocab, idx2token
        
    def __len__(self):
        return len(self.tokens) // self.seq_len

    def __getitem__(self, idx):
        chunk = self.tokens[idx * self.seq_len: (idx + 1) * self.seq_len]
        x = chunk[:-1]
        y = chunk[1:]
        return torch.tensor(x), torch.tensor(y)

def collate_fn(batch):
    inputs, targets = zip(*batch)
    inputs = pad_sequence(inputs, batch_first=True, padding_value=0)
    targets = pad_sequence(targets, batch_first=True, padding_value=0)
    return inputs, targets

# Set the path to your text file and define sequence length
path_to_text = 'C:/Users/Dell-PC/Desktop/Simple_Transfomer.txt'  # replace with the path to your text file
seq_len = 30  # sequence length

# Create a dataset and data loader
dataset = TextDataset(path_to_text, seq_len)
data_loader = DataLoader(dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)

class SingleLSTM_TXT(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers=3):
        # In a language model, input_size is usually the dimension of the embedding vector,
        # and output_size is the size of the vocabulary.
        super(SingleLSTM_TXT, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=output_size, embedding_dim=input_size)
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        embedded = self.embedding(x)
        h, _ = self.lstm(embedded)
        # You may need to reshape h if you want to include more than the last LSTM output
        h = h.reshape(-1, h.shape[-1])
        return self.fc(h)
    
def train_model_TXT(model, criterion, optimizer, num_epochs, data_loader):
    model.train()
    for epoch in range(num_epochs):
        total_loss = 0

        for inputs, targets in data_loader:
            optimizer.zero_grad()
            predictions = model(inputs)
            loss = criterion(predictions, targets.view(-1))
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        average_loss = total_loss / len(data_loader.dataset)
        print(f"Epoch {epoch+1}, Loss: {average_loss}")

def generate_text(model, dataset, seed_text, num_generate, temperature=1.0):
    model.eval()  # Put the model in evaluation mode
    
    # List to store the generated tokens
    generated_tokens = []

    # Initial sequence (prefix) to start the generation process
    input_sequence = [dataset.vocab.get(word, 0) for word in seed_text.split()]  # Convert to token IDs
    current_sequence = torch.tensor(input_sequence, dtype=torch.long).unsqueeze(0)

    # Generate num_generate tokens
    for _ in range(num_generate):
        # Forward pass through the model
        with torch.no_grad():
            output = model(current_sequence)
        
        # Get probabilities, apply temperature scaling, and sample from the distribution
        probabilities = F.softmax(output[-1] / temperature, dim=0).detach()
        next_token_idx = torch.multinomial(probabilities, 1).item()

        # Append token to the current sequence and to the generated tokens
        generated_tokens.append(next_token_idx)
        current_sequence = torch.cat((current_sequence, torch.tensor([[next_token_idx]])), 1)
    
    # Convert tokens to words
    generated_text = " ".join([dataset.idx2token[token] for token in generated_tokens])
    return generated_text

# Instantiate the SingleLSTM_TXT model
vocab_size = len(dataset.vocab)  # The size of the vocabulary
embedding_dim = 128  # Size of embeddings
hidden_size = 128  # Number of features in the hidden state of the LSTM
num_layers = 4  # Number of stacked LSTM layers

single_lstm_txt_model = SingleLSTM_TXT(vocab_size, embedding_dim, hidden_size, num_layers)

# Training parameters
num_epochs = 10
batch_size = 128
learning_rate = 1e-4

# Define Loss Function and Optimizer for TXT dataset
criterion = nn.CrossEntropyLoss()  # Use CrossEntropyLoss for classification tasks
optimizer = torch.optim.Adam(single_lstm_txt_model.parameters(), lr=learning_rate)

# Train the model with the text data
train_model_TXT(single_lstm_txt_model, criterion, optimizer, num_epochs, data_loader)

# Seed text to start the generation
seed_text = "Give me code to define a "
num_generate = 8  # Number of words to generate
temperature = 1.0  # Sampling temperature, higher will increase diversity

# Use the trained model to generate text
generated_text = generate_text(single_lstm_txt_model, dataset, seed_text, num_generate, temperature)
print("Generated Text: ", generated_text)