| import os
|
| os.environ["OMP_NUM_THREADS"] = "1"
|
| os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
| import streamlit as st
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from sentence_transformers import SentenceTransformer
|
| from sklearn.feature_extraction.text import CountVectorizer
|
| import pandas as pd
|
| import numpy as np
|
| import time
|
|
|
|
|
|
|
|
|
| st.set_page_config(page_title="Neural VAE Architecture", layout="wide")
|
|
|
|
|
| st.markdown("""
|
| <style>
|
| html { overflow-y: scroll; }
|
| .block-container { overflow-x: hidden; }
|
| </style>
|
| """, unsafe_allow_html=True)
|
|
|
| st.title("Variational Autoencoder (VAE) - Topic Modeling with Math Visualization")
|
| st.markdown("Visualize the Neural Network transform text into topics in real-time.")
|
|
|
|
|
| try:
|
| st.image("pipeline.png", use_container_width=True)
|
| except Exception:
|
| st.info("💡 Tip: Upload a 'pipeline.png' image to display the architecture diagram here.")
|
|
|
| st.divider()
|
|
|
|
|
| @st.cache_resource
|
| def load_bert():
|
| return SentenceTransformer("all-MiniLM-L6-v2")
|
|
|
|
|
| st.sidebar.header("1. Input & Text Prep")
|
| n_topics = st.sidebar.slider("Number of Topics (Latent Dimensions)", 2, 10, 3)
|
| words_per_theme = st.sidebar.slider("Top Words per Topic", 3, 10, 5)
|
|
|
| st.sidebar.header("2. Neural Network Settings")
|
| activation_choice = st.sidebar.selectbox("Encoder Activation", ["ReLU", "Softplus", "GELU"])
|
| dropout_rate = st.sidebar.slider("Dropout Rate (Noise)", 0.0, 0.5, 0.2)
|
|
|
| st.sidebar.header("3. Optimization (Loss)")
|
| kl_beta = st.sidebar.slider("Max KL Beta (β)", 0.0, 1.0, 0.1)
|
| kl_warmup = st.sidebar.slider("KL Annealing (Epochs)", 0, 500, 200)
|
| learning_rate = st.sidebar.selectbox("Learning Rate", [0.001, 0.005, 0.01, 0.05], index=2)
|
| epochs = st.sidebar.slider("Training Epochs", 100, 1000, 500)
|
|
|
| sample_abstract = """
|
| Students who registered for the Mapping with Google massive open online course (MOOC)
|
| were asked several questions during the registration process to identify prior
|
| experience with eleven skills as well as their goals for registering for the course.
|
| At the end of the course, we compared students' self reports of goal achievement
|
| with behavioral click-stream analysis. In addition, we assessed how well prior
|
| skill in a subject predicts a student's course completion and found no correlation.
|
| """
|
| raw_data = st.sidebar.text_area("Training Data:", value=sample_abstract, height=150)
|
|
|
|
|
|
|
|
|
| class ContextualVAE(nn.Module):
|
| def __init__(self, input_dim, hidden_dim, latent_dim, vocab_size, activation, dropout):
|
| super().__init__()
|
| act_fn = nn.ReLU() if activation == "ReLU" else nn.Softplus() if activation == "Softplus" else nn.GELU()
|
|
|
| self.encoder = nn.Sequential(
|
| nn.Linear(input_dim, hidden_dim),
|
| act_fn,
|
| nn.Dropout(dropout),
|
| nn.BatchNorm1d(hidden_dim)
|
| )
|
| self.fc_mu = nn.Linear(hidden_dim, latent_dim)
|
| self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
|
| self.decoder = nn.Linear(latent_dim, vocab_size, bias=False)
|
|
|
| def encode(self, x):
|
| h = self.encoder(x)
|
| return self.fc_mu(h), self.fc_logvar(h)
|
|
|
| def reparameterize(self, mu, logvar):
|
| std = torch.exp(0.5 * logvar)
|
| eps = torch.randn_like(std)
|
| return mu + eps * std
|
|
|
| def forward(self, x):
|
| mu, logvar = self.encode(x)
|
| z = self.reparameterize(mu, logvar)
|
| p_topic = F.softmax(z, dim=-1)
|
| recon_log_probs = F.log_softmax(self.decoder(p_topic), dim=-1)
|
| return recon_log_probs, mu, logvar, z
|
|
|
|
|
|
|
|
|
| st.divider()
|
|
|
| col1, col2, col3 = st.columns(3)
|
|
|
| with col1:
|
| st.markdown("### 📥 1. The Encoder")
|
| st.latex(r"\mu, \log(\sigma^2) = \text{NeuralNet}(X_{bert})")
|
| st.caption("Compresses 384D BERT vectors into Mean and Variance coordinates.")
|
| ui_mu = st.empty()
|
| ui_var = st.empty()
|
| with st.expander("View Encoder History"):
|
| ui_encoder_table = st.empty()
|
|
|
| with col2:
|
| st.markdown("### 🕳️ 2. The Latent Space")
|
| st.latex(r"Z = \mu + \sigma \odot \epsilon")
|
| st.caption("The Reparameterization Trick. Samples a topic vector adding random noise.")
|
| ui_z = st.empty()
|
| ui_softmax = st.empty()
|
| with st.expander("View Latent History"):
|
| ui_latent_table = st.empty()
|
|
|
| with col3:
|
| st.markdown("### 📤 3. The Decoder")
|
| st.latex(r"\hat{Y} = \text{Softmax}(Z) \cdot W_{decoder}")
|
| st.caption("Maps the tiny topic vector back to the full Vocabulary size.")
|
| ui_decoder = st.empty()
|
| with st.expander("View Decoder History"):
|
| ui_decoder_table = st.empty()
|
|
|
| st.divider()
|
| st.markdown("### 📉 4. Optimization (Backpropagation)")
|
| st.latex(r"\mathcal{L} = D_{KL}(Y || \hat{Y}) + \beta \cdot D_{KL}(\mathcal{N}(\mu, \sigma) || \mathcal{N}(0, I))")
|
|
|
| col_loss_text, col_loss_graph = st.columns([1, 2])
|
| with col_loss_text:
|
| ui_beta = st.empty()
|
| ui_recon_loss = st.empty()
|
| ui_kl_loss = st.empty()
|
| with col_loss_graph:
|
| loss_chart_placeholder = st.empty()
|
|
|
|
|
|
|
|
|
| if st.sidebar.button("🚀 Start Training and Backpropagation", type="primary", use_container_width=True):
|
|
|
| sentences = [s.strip() for s in raw_data.split('.') if len(s.strip()) > 10]
|
|
|
| stop_w = list(CountVectorizer(stop_words='english').get_stop_words()) + ['students', 'course']
|
| vectorizer = CountVectorizer(stop_words=stop_w)
|
| bow_matrix = vectorizer.fit_transform(sentences).toarray()
|
| vocab = vectorizer.get_feature_names_out()
|
|
|
| target_dist = torch.tensor(bow_matrix).float()
|
| target_dist = target_dist / (target_dist.sum(dim=1, keepdim=True) + 1e-10)
|
|
|
| with st.spinner("Loading BERT (Cached)..."):
|
| embedder = load_bert()
|
| bert_inputs = torch.tensor(embedder.encode(sentences))
|
|
|
| model = ContextualVAE(384, 128, n_topics, len(vocab), activation_choice, dropout_rate)
|
| optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
|
|
| recon_hist, kl_hist = [], []
|
| history_encoder, history_latent, history_decoder = [], [], []
|
|
|
|
|
| for epoch in range(1, epochs + 1):
|
| model.train()
|
| optimizer.zero_grad()
|
|
|
| recon, mu, logvar, z = model(bert_inputs)
|
|
|
| recon_loss = F.kl_div(recon, target_dist, reduction='batchmean')
|
| kl_loss = -0.5 * torch.mean(torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1))
|
|
|
| current_beta = kl_beta * min(1.0, epoch / max(1, kl_warmup))
|
| total_loss = recon_loss + (current_beta * kl_loss)
|
|
|
| total_loss.backward()
|
| optimizer.step()
|
|
|
|
|
| s1_mu = np.round(mu[0].detach().numpy(), 3).tolist()
|
| s1_var = np.round(torch.exp(0.5 * logvar[0]).detach().numpy(), 3).tolist()
|
| s1_z = np.round(z[0].detach().numpy(), 3).tolist()
|
| s1_soft = np.round(F.softmax(z[0], dim=-1).detach().numpy(), 3).tolist()
|
|
|
| history_encoder.append({"Epoch": epoch, "Mean (μ)": str(s1_mu), "Variance (σ)": str(s1_var)})
|
| history_latent.append({"Epoch": epoch, "Sampled Vector (Z)": str(s1_z), "Topic Probs": str(s1_soft)})
|
|
|
| weights = model.decoder.weight.T
|
| dec_dict = {"Epoch": epoch}
|
| dec_text = ""
|
|
|
| for i in range(n_topics):
|
| top_idx = torch.argsort(weights[i], descending=True)[:words_per_theme]
|
| theme_words = [vocab[j] for j in top_idx]
|
| dec_text += f"**T{i+1}:** {', '.join(theme_words)}\n\n"
|
| dec_dict[f"Topic {i+1}"] = ", ".join(theme_words)
|
|
|
| history_decoder.append(dec_dict)
|
|
|
| recon_hist.append(recon_loss.item())
|
| kl_hist.append((current_beta * kl_loss).item())
|
|
|
|
|
|
|
|
|
| if epoch % 20 == 0 or epoch == epochs:
|
|
|
| ui_mu.info(f"**Mean (μ) for Sentence 1:**\n`{s1_mu}`")
|
| ui_var.warning(f"**Variance (σ) for Sentence 1:**\n`{s1_var}`")
|
| ui_z.success(f"**Sampled 'Z' for Sentence 1:**\n`{s1_z}`")
|
| ui_softmax.error(f"**Topic Probabilities:**\n`{s1_soft}`")
|
| ui_decoder.info(dec_text)
|
|
|
|
|
| ui_encoder_table.dataframe(pd.DataFrame(history_encoder), hide_index=True, use_container_width=True, height=200)
|
| ui_latent_table.dataframe(pd.DataFrame(history_latent), hide_index=True, use_container_width=True, height=200)
|
| ui_decoder_table.dataframe(pd.DataFrame(history_decoder), hide_index=True, use_container_width=True, height=200)
|
|
|
| ui_beta.metric("Current Beta (Annealing)", f"{current_beta:.4f}")
|
| ui_recon_loss.metric("Reconstruction Loss", f"{recon_loss.item():.4f}")
|
| ui_kl_loss.metric("KL Penalty", f"{(current_beta * kl_loss).item():.4f}")
|
|
|
| chart_df = pd.DataFrame({"Recon Loss": recon_hist, "KL Penalty": kl_hist})
|
| loss_chart_placeholder.line_chart(chart_df)
|
|
|
| time.sleep(0.01)
|
|
|
| st.success("Neural Network Training Complete!") |
