Update src/streamlit_app.py

src/streamlit_app.py

@@ -38,9 +38,11 @@ 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)
 
+# Number of tokens from dataset to use for training (minimum 100 tokens)
+num_train_tokens = st.sidebar.slider("Number of tokens from dataset to train on", min_value=100, max_value=100000, value=1000, step=100)
+
 train_button = st.sidebar.button("Train Model")
 
 device = torch.device("cpu")  # force CPU usage
@@ -74,10 +76,12 @@ def tokenize(text, tokenizer_type):
         tokens = text.split()
     return tokens
 
-tokens = tokenize(text_data, tokenizer_type)
-vocab = list(set(tokens))
+tokens_all = tokenize(text_data, tokenizer_type)
+
+# Cap tokens to requested number for training
+tokens = tokens_all[:num_train_tokens]
 
-# Add PAD token to vocab for padding contexts shorter than context_size - 1
+vocab = list(set(tokens))
 PAD_TOKEN = "<PAD>"
 if PAD_TOKEN not in vocab:
     vocab.append(PAD_TOKEN)
@@ -90,10 +94,6 @@ idx_to_token = {i: tok for tok, i in token_to_idx.items()}
 ###################################
 
 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
@@ -145,31 +145,33 @@ class FFNN(nn.Module):
 
 def train_ffnn(tokens, context_size=3, epochs=3):
     data = []
-    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]
+    for i in range(len(tokens) - (context_size - 1)):
+        context = tokens[i : i + context_size - 1]
         context = pad_context(context, context_size - 1)
-        target = tokens[i]
+        target = tokens[i + context_size - 1]
         data.append((
             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])
         ))
 
+    if len(data) == 0:
+        st.warning("No training data generated. Increase dataset size or reduce context size.")
+        return None
+
     model = FFNN(len(vocab), context_size - 1).to(device)
     optimizer = optim.Adam(model.parameters(), lr=0.01)
     criterion = nn.CrossEntropyLoss()
 
     progress_bar = st.progress(0)
-    total_steps = epochs * len(data)
+    total_steps = len(data) * epochs
     step = 0
 
     model.train()
-
     for epoch in range(epochs):
         total_loss = 0
         random.shuffle(data)
         for x, y in data:
-            x = x.unsqueeze(0)  # batch size 1
+            x = x.unsqueeze(0)
             y = torch.tensor([y], device=device)
 
             optimizer.zero_grad()
@@ -201,11 +203,10 @@ def ffnn_predict(model, context, temperature=1.0):
 
 def train_dt(tokens, context_size=3):
     X, y = [], []
-    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]
+    for i in range(len(tokens) - (context_size - 1)):
+        context = tokens[i : i + context_size - 1]
         context = pad_context(context, context_size - 1)
-        target = tokens[i]
+        target = tokens[i + context_size - 1]
         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]))
 
@@ -226,11 +227,10 @@ def dt_predict(model, context):
 
 def train_gbt(tokens, context_size=3):
     X, y = [], []
-    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]
+    for i in range(len(tokens) - (context_size - 1)):
+        context = tokens[i : i + context_size - 1]
         context = pad_context(context, context_size - 1)
-        target = tokens[i]
+        target = tokens[i + context_size - 1]
         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]))
 
@@ -264,26 +264,28 @@ class RNNModel(nn.Module):
 
 def train_rnn(tokens, context_size=3, epochs=3):
     data = []
-    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]
+    for i in range(len(tokens) - (context_size - 1)):
+        context = tokens[i : i + context_size - 1]
         context = pad_context(context, context_size - 1)
-        target = tokens[i]
+        target = tokens[i + context_size - 1]
         data.append((
             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])
         ))
 
+    if len(data) == 0:
+        st.warning("No training data generated. Increase dataset size or reduce context size.")
+        return None
+
     model = RNNModel(len(vocab)).to(device)
     optimizer = optim.Adam(model.parameters(), lr=0.01)
     criterion = nn.CrossEntropyLoss()
 
     progress_bar = st.progress(0)
-    total_steps = epochs * len(data)
+    total_steps = len(data) * epochs
     step = 0
 
     model.train()
-
     for epoch in range(epochs):
         total_loss = 0
         h = None
@@ -296,8 +298,8 @@ def train_rnn(tokens, context_size=3, epochs=3):
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
-            total_loss += loss.item()
 
+            total_loss += loss.item()
             step += 1
             progress_bar.progress(step / total_steps)
 
@@ -319,7 +321,7 @@ def rnn_predict(model, context, temperature=1.0):
 ###################################
 
 if train_button:
-    st.write(f"Training **{model_type}** model with context size {context_size}...")
+    st.write(f"Training **{model_type}** model with context size {context_size} on {len(tokens)} tokens...")
 
     if model_type == "N-gram":
         with st.spinner("Training N-gram model..."):
@@ -333,10 +335,13 @@ if train_button:
     elif model_type == "RNN":
         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.")
+    if model is not None:
+        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.")
+    else:
+        st.error("Training failed due to no data.")
 
 ###################################
 # Chat interface