Spaces:
Running
Running
File size: 4,025 Bytes
8462f4c |
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 |
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.float32).unsqueeze(0).unsqueeze(-1)
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.float32).unsqueeze(0).unsqueeze(-1)
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
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])
X = np.reshape(dataX, (len(dataX), seq_length, 1))
X = X / float(len(chars))
Y = np.eye(len(chars))[dataY]
# Convert to PyTorch tensors
X_tensor = torch.tensor(X, dtype=torch.float32)
Y_tensor = torch.tensor(dataY, 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, X_tensor.size(0), hidden_size), torch.zeros(num_layers, X_tensor.size(0), hidden_size))
for i in range(0, X_tensor.size(0), seq_length):
inputs = X_tensor[i:i + seq_length]
targets = Y_tensor[i:i + seq_length]
# 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()
st.write(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.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)
|