Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,3 @@
-# app.py
-
 import os
 
 # ✅ Fix PermissionError on Hugging Face Spaces
@@ -39,6 +37,10 @@ model_type = st.sidebar.selectbox(
 )
 
 temperature = st.sidebar.slider("Sampling Temperature", 0.1, 2.0, 1.0)
+
+# Context size slider (minimum 2)
+context_size = st.sidebar.slider("Context Size (how many tokens to look back)", min_value=2, max_value=10, value=3, step=1)
+
 train_button = st.sidebar.button("Train Model")
 
 device = torch.device("cpu")  # force CPU usage
@@ -74,9 +76,30 @@ def tokenize(text, tokenizer_type):
 
 tokens = tokenize(text_data, tokenizer_type)
 vocab = list(set(tokens))
+
+# Add PAD token to vocab for padding contexts shorter than context_size - 1
+PAD_TOKEN = "<PAD>"
+if PAD_TOKEN not in vocab:
+    vocab.append(PAD_TOKEN)
+
 token_to_idx = {tok: i for i, tok in enumerate(vocab)}
 idx_to_token = {i: tok for tok, i in token_to_idx.items()}
 
+###################################
+# Helper to pad context
+###################################
+
+def pad_context(context, size):
+    """
+    Pads the context list at the front with PAD_TOKEN if length < size,
+    or truncates to last `size` tokens if longer.
+    """
+    pad_len = size - len(context)
+    if pad_len > 0:
+        return [PAD_TOKEN]*pad_len + context
+    else:
+        return context[-size:]
+
 ###################################
 # Models
 ###################################
@@ -122,15 +145,17 @@ class FFNN(nn.Module):
 
 def train_ffnn(tokens, context_size=3, epochs=3):
     data = []
-    for i in range(len(tokens) - context_size):
-        context = tokens[i:i+context_size-1]
-        target = tokens[i+context_size-1]
+    for i in range(len(tokens)):
+        start_idx = i - (context_size - 1)
+        context = tokens[start_idx:i] if start_idx >= 0 else tokens[0:i]
+        context = pad_context(context, context_size - 1)
+        target = tokens[i]
         data.append((
-            torch.tensor([token_to_idx[tok] for tok in context], device=device),
-            token_to_idx[target]
+            torch.tensor([token_to_idx.get(t, token_to_idx[PAD_TOKEN]) for t in context], device=device),
+            token_to_idx.get(target, token_to_idx[PAD_TOKEN])
         ))
 
-    model = FFNN(len(vocab), context_size-1).to(device)
+    model = FFNN(len(vocab), context_size - 1).to(device)
     optimizer = optim.Adam(model.parameters(), lr=0.01)
     criterion = nn.CrossEntropyLoss()
 
@@ -138,28 +163,33 @@ def train_ffnn(tokens, context_size=3, epochs=3):
     total_steps = epochs * len(data)
     step = 0
 
+    model.train()
+
     for epoch in range(epochs):
         total_loss = 0
+        random.shuffle(data)
         for x, y in data:
-            x = x.unsqueeze(0)
+            x = x.unsqueeze(0)  # batch size 1
             y = torch.tensor([y], device=device)
+
+            optimizer.zero_grad()
             out = model(x)
             loss = criterion(out, y)
-            optimizer.zero_grad()
             loss.backward()
             optimizer.step()
-            total_loss += loss.item()
 
+            total_loss += loss.item()
             step += 1
             progress_bar.progress(step / total_steps)
 
-        st.write(f"Epoch {epoch+1}, Loss: {total_loss:.4f}")
+        st.write(f"Epoch {epoch+1}, Loss: {total_loss/len(data):.4f}")
 
     progress_bar.empty()
     return model
 
 def ffnn_predict(model, context, temperature=1.0):
-    x = torch.tensor([token_to_idx.get(tok, 0) for tok in context[-2:]], device=device).unsqueeze(0)
+    context = pad_context(context, context_size - 1)
+    x = torch.tensor([token_to_idx.get(tok, token_to_idx[PAD_TOKEN]) for tok in context], device=device).unsqueeze(0)
     with torch.no_grad():
         logits = model(x).squeeze()
         probs = torch.softmax(logits / temperature, dim=0).cpu().numpy()
@@ -171,11 +201,13 @@ def ffnn_predict(model, context, temperature=1.0):
 
 def train_dt(tokens, context_size=3):
     X, y = [], []
-    for i in range(len(tokens) - context_size):
-        context = tokens[i:i+context_size-1]
-        target = tokens[i+context_size-1]
-        X.append([token_to_idx[tok] for tok in context])
-        y.append(token_to_idx[target])
+    for i in range(len(tokens)):
+        start_idx = i - (context_size - 1)
+        context = tokens[start_idx:i] if start_idx >= 0 else tokens[0:i]
+        context = pad_context(context, context_size - 1)
+        target = tokens[i]
+        X.append([token_to_idx.get(t, token_to_idx[PAD_TOKEN]) for t in context])
+        y.append(token_to_idx.get(target, token_to_idx[PAD_TOKEN]))
 
     with st.spinner("Training Decision Tree..."):
         model = DecisionTreeClassifier()
@@ -183,7 +215,8 @@ def train_dt(tokens, context_size=3):
     return model
 
 def dt_predict(model, context):
-    x = [token_to_idx.get(tok, 0) for tok in context[-2:]]
+    context = pad_context(context, context_size - 1)
+    x = [token_to_idx.get(tok, token_to_idx[PAD_TOKEN]) for tok in context]
     pred = model.predict([x])[0]
     return idx_to_token[pred]
 
@@ -193,11 +226,13 @@ def dt_predict(model, context):
 
 def train_gbt(tokens, context_size=3):
     X, y = [], []
-    for i in range(len(tokens) - context_size):
-        context = tokens[i:i+context_size-1]
-        target = tokens[i+context_size-1]
-        X.append([token_to_idx[tok] for tok in context])
-        y.append(token_to_idx[target])
+    for i in range(len(tokens)):
+        start_idx = i - (context_size - 1)
+        context = tokens[start_idx:i] if start_idx >= 0 else tokens[0:i]
+        context = pad_context(context, context_size - 1)
+        target = tokens[i]
+        X.append([token_to_idx.get(t, token_to_idx[PAD_TOKEN]) for t in context])
+        y.append(token_to_idx.get(target, token_to_idx[PAD_TOKEN]))
 
     with st.spinner("Training Gradient Boosted Tree..."):
         model = GradientBoostingClassifier()
@@ -205,7 +240,8 @@ def train_gbt(tokens, context_size=3):
     return model
 
 def gbt_predict(model, context):
-    x = [token_to_idx.get(tok, 0) for tok in context[-2:]]
+    context = pad_context(context, context_size - 1)
+    x = [token_to_idx.get(tok, token_to_idx[PAD_TOKEN]) for tok in context]
     pred = model.predict([x])[0]
     return idx_to_token[pred]
 
@@ -228,12 +264,14 @@ class RNNModel(nn.Module):
 
 def train_rnn(tokens, context_size=3, epochs=3):
     data = []
-    for i in range(len(tokens) - context_size):
-        context = tokens[i:i+context_size-1]
-        target = tokens[i+context_size-1]
+    for i in range(len(tokens)):
+        start_idx = i - (context_size - 1)
+        context = tokens[start_idx:i] if start_idx >= 0 else tokens[0:i]
+        context = pad_context(context, context_size - 1)
+        target = tokens[i]
         data.append((
-            torch.tensor([token_to_idx[tok] for tok in context], device=device),
-            token_to_idx[target]
+            torch.tensor([token_to_idx.get(t, token_to_idx[PAD_TOKEN]) for t in context], device=device),
+            token_to_idx.get(target, token_to_idx[PAD_TOKEN])
         ))
 
     model = RNNModel(len(vocab)).to(device)
@@ -244,9 +282,12 @@ def train_rnn(tokens, context_size=3, epochs=3):
     total_steps = epochs * len(data)
     step = 0
 
+    model.train()
+
     for epoch in range(epochs):
         total_loss = 0
         h = None
+        random.shuffle(data)
         for x, y in data:
             x = x.unsqueeze(0)
             y = torch.tensor([y], device=device)
@@ -260,13 +301,14 @@ def train_rnn(tokens, context_size=3, epochs=3):
             step += 1
             progress_bar.progress(step / total_steps)
 
-        st.write(f"Epoch {epoch+1}, Loss: {total_loss:.4f}")
+        st.write(f"Epoch {epoch+1}, Loss: {total_loss/len(data):.4f}")
 
     progress_bar.empty()
     return model
 
 def rnn_predict(model, context, temperature=1.0):
-    x = torch.tensor([token_to_idx.get(tok, 0) for tok in context[-2:]], device=device).unsqueeze(0)
+    context = pad_context(context, context_size - 1)
+    x = torch.tensor([token_to_idx.get(tok, token_to_idx[PAD_TOKEN]) for tok in context], device=device).unsqueeze(0)
     with torch.no_grad():
         logits, _ = model(x)
         probs = torch.softmax(logits.squeeze() / temperature, dim=0).cpu().numpy()
@@ -277,22 +319,23 @@ def rnn_predict(model, context, temperature=1.0):
 ###################################
 
 if train_button:
-    st.write(f"Training **{model_type}** model...")
+    st.write(f"Training **{model_type}** model with context size {context_size}...")
 
     if model_type == "N-gram":
         with st.spinner("Training N-gram model..."):
-            model = NGramModel(tokens, n=3)
+            model = NGramModel(tokens, n=context_size)
     elif model_type == "Feed Forward NN":
-        model = train_ffnn(tokens)
+        model = train_ffnn(tokens, context_size=context_size)
     elif model_type == "Decision Tree":
-        model = train_dt(tokens)
+        model = train_dt(tokens, context_size=context_size)
     elif model_type == "Gradient Boosted Tree":
-        model = train_gbt(tokens)
+        model = train_gbt(tokens, context_size=context_size)
     elif model_type == "RNN":
-        model = train_rnn(tokens)
+        model = train_rnn(tokens, context_size=context_size)
 
     st.session_state["model"] = model
     st.session_state["model_type"] = model_type
+    st.session_state["context_size"] = context_size
     st.success(f"{model_type} model trained.")
 
 ###################################
@@ -309,16 +352,18 @@ if "model" in st.session_state:
         generated = context.copy()
 
         for _ in range(20):
+            ctx = pad_context(generated, st.session_state["context_size"] - 1)
+
             if st.session_state["model_type"] == "N-gram":
-                next_tok = st.session_state["model"].predict(generated, temperature)
+                next_tok = st.session_state["model"].predict(ctx, temperature)
             elif st.session_state["model_type"] == "Feed Forward NN":
-                next_tok = ffnn_predict(st.session_state["model"], generated, temperature)
+                next_tok = ffnn_predict(st.session_state["model"], ctx, temperature)
             elif st.session_state["model_type"] == "Decision Tree":
-                next_tok = dt_predict(st.session_state["model"], generated)
+                next_tok = dt_predict(st.session_state["model"], ctx)
             elif st.session_state["model_type"] == "Gradient Boosted Tree":
-                next_tok = gbt_predict(st.session_state["model"], generated)
+                next_tok = gbt_predict(st.session_state["model"], ctx)
             elif st.session_state["model_type"] == "RNN":
-                next_tok = rnn_predict(st.session_state["model"], generated, temperature)
+                next_tok = rnn_predict(st.session_state["model"], ctx, temperature)
 
             generated.append(next_tok)
             if next_tok == "<END>":