File size: 4,618 Bytes
8462f4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b980d33
8462f4c
 
 
 
 
 
 
 
b980d33
8462f4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b8c5c1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b980d33
b8c5c1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b980d33
b8c5c1a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import streamlit as st
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# Define the RNN or LSTM Model
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        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, h):
        out, h = self.lstm(x, h)
        out = self.fc(out[:, -1, :])
        return out, h

# Text generation function
def generate_text(model, start_str, length, char_to_int, int_to_char, num_layers, hidden_size):
    model.eval()
    input_seq = [char_to_int[c] for c in start_str]
    input_seq = torch.tensor(input_seq, dtype=torch.long).unsqueeze(0)
    h = (torch.zeros(num_layers, 1, hidden_size), torch.zeros(num_layers, 1, hidden_size))
    generated_text = start_str

    for _ in range(length):
        output, h = model(input_seq, h)
        _, predicted = torch.max(output, 1)
        predicted_char = int_to_char[predicted.item()]
        generated_text += predicted_char
        input_seq = torch.tensor([[char_to_int[predicted_char]]], dtype=torch.long)

    return generated_text

# Streamlit interface
st.title("RNN/LSTM Text Generation")

# Inputs
text_data = st.text_area("Enter your text data for training:")
start_string = st.text_input("Enter the start string for text generation:")
seq_length = st.number_input("Sequence length:", min_value=10, value=100)
hidden_size = st.number_input("Hidden size:", min_value=50, value=256)
num_layers = st.number_input("Number of layers:", min_value=1, value=2)
learning_rate = st.number_input("Learning rate:", min_value=0.0001, value=0.003, format="%.4f")
num_epochs = st.number_input("Number of epochs:", min_value=1, value=20)
generate_length = st.number_input("Generated text length:", min_value=50, value=500)

if st.button("Train and Generate"):
    # Data Preparation
    text = text_data
    if len(text) <= seq_length:
        st.error("Text data is too short for the given sequence length. Please enter more text data.")
    else:
        chars = sorted(list(set(text)))
        char_to_int = {c: i for i, c in enumerate(chars)}
        int_to_char = {i: c for i, c in enumerate(chars)}

        # Prepare input-output pairs
        dataX = []
        dataY = []
        for i in range(0, len(text) - seq_length):
            seq_in = text[i:i + seq_length]
            seq_out = text[i + seq_length]
            dataX.append([char_to_int[char] for char in seq_in])
            dataY.append(char_to_int[seq_out])

        if len(dataX) == 0:
            st.error("Not enough data to create input-output pairs. Please provide more text data.")
        else:
            X = np.reshape(dataX, (len(dataX), seq_length))
            X = X / float(len(chars))
            Y = np.array(dataY)

            # Convert to PyTorch tensors
            X_tensor = torch.tensor(X, dtype=torch.float32)
            Y_tensor = torch.tensor(Y, dtype=torch.long)

            # Model initialization
            model = LSTMModel(input_size=1, hidden_size=hidden_size, output_size=len(chars), num_layers=num_layers)

            # Loss and optimizer
            criterion = nn.CrossEntropyLoss()
            optimizer = optim.Adam(model.parameters(), lr=learning_rate)

            # Training the model
            for epoch in range(num_epochs):
                h = (torch.zeros(num_layers, 1, hidden_size), torch.zeros(num_layers, 1, hidden_size))
                epoch_loss = 0
                for i in range(len(dataX)):
                    inputs = X_tensor[i].unsqueeze(0).unsqueeze(-1)
                    targets = Y_tensor[i].unsqueeze(0)

                    # Forward pass
                    outputs, h = model(inputs, h)
                    h = (h[0].detach(), h[1].detach())
                    loss = criterion(outputs, targets)

                    # Backward pass and optimization
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()

                    epoch_loss += loss.item()

                avg_loss = epoch_loss / len(dataX)
                st.write(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}')

            # Text generation
            generated_text = generate_text(model, start_string, generate_length, char_to_int, int_to_char, num_layers, hidden_size)
            st.subheader("Generated Text")
            st.write(generated_text)