Blame for app.py

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

# ==========================================
# 1. PAGE CONFIG & UI SETUP (WITH ANTI-GLITCH CSS)
# ==========================================
st.set_page_config(page_title="Neural VAE Architecture", layout="wide")

# FIX: CSS injection to permanently lock the scrollbar and prevent horizontal jumping!
st.markdown("""

    <style>

        html { overflow-y: scroll; }

        .block-container { overflow-x: hidden; }

    </style>

    """, unsafe_allow_html=True)

st.title("🧠 Variational Autoencoder (VAE): Inside the Black Box")
st.markdown("Watch the mathematics of a Neural Network transform text into topics in real-time.")

# --- DISPLAY ARCHITECTURE DIAGRAM ---
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()

# --- CACHE THE HEAVY BERT MODEL ---
@st.cache_resource
def load_bert():
    return SentenceTransformer("all-MiniLM-L6-v2")

# --- SIDEBAR CONTROLS ---
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)

# ==========================================
# 2. THE PYTORCH VAE ARCHITECTURE
# ==========================================
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

# ==========================================
# 3. THE ARCHITECTURE DASHBOARD (UI LAYOUT)
# ==========================================
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()

# ==========================================
# 4. TRAINING EXECUTION
# ==========================================
if st.sidebar.button("🚀 Start 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 = [], [], []
    
# --- PYTORCH TRAINING LOOP ---
    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()
        
        # Grab math for Sentence 1
        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())
        
        # ==========================================
        # 2. SAFE UI RENDER (Every 20 epochs)
        # ==========================================
        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)
            
            # FIX: Added `height=200` to prevent the dataframe from stretching the webpage!
            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) # Slightly faster sleep to prevent HF WebSocket timeouts

    st.success("Neural Network Training Complete!")