Update src/streamlit_app.py

src/streamlit_app.py

@@ -159,3 +159,177 @@ def train_ffnn(tokens, context_size=3, epochs=3):
     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)
+    with torch.no_grad():
+        logits = model(x).squeeze()
+        probs = torch.softmax(logits / temperature, dim=0).cpu().numpy()
+        return np.random.choice(vocab, p=probs)
+
+###################################
+# Decision Tree
+###################################
+
+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])
+
+    with st.spinner("Training Decision Tree..."):
+        model = DecisionTreeClassifier()
+        model.fit(X, y)
+    return model
+
+def dt_predict(model, context):
+    x = [token_to_idx.get(tok, 0) for tok in context[-2:]]
+    pred = model.predict([x])[0]
+    return idx_to_token[pred]
+
+###################################
+# Gradient Boosted Tree
+###################################
+
+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])
+
+    with st.spinner("Training Gradient Boosted Tree..."):
+        model = GradientBoostingClassifier()
+        model.fit(X, y)
+    return model
+
+def gbt_predict(model, context):
+    x = [token_to_idx.get(tok, 0) for tok in context[-2:]]
+    pred = model.predict([x])[0]
+    return idx_to_token[pred]
+
+###################################
+# RNN
+###################################
+
+class RNNModel(nn.Module):
+    def __init__(self, vocab_size, embed_size=64, hidden_size=128):
+        super().__init__()
+        self.embed = nn.Embedding(vocab_size, embed_size)
+        self.rnn = nn.RNN(embed_size, hidden_size, batch_first=True)
+        self.fc = nn.Linear(hidden_size, vocab_size)
+
+    def forward(self, x, h=None):
+        x = self.embed(x)
+        out, h = self.rnn(x, h)
+        out = self.fc(out[:, -1, :])
+        return out, h
+
+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]
+        data.append((
+            torch.tensor([token_to_idx[tok] for tok in context], device=device),
+            token_to_idx[target]
+        ))
+
+    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)
+    step = 0
+
+    for epoch in range(epochs):
+        total_loss = 0
+        h = None
+        for x, y in data:
+            x = x.unsqueeze(0)
+            y = torch.tensor([y], device=device)
+            out, h = model(x, h)
+            loss = criterion(out, y)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+
+            step += 1
+            progress_bar.progress(step / total_steps)
+
+        st.write(f"Epoch {epoch+1}, Loss: {total_loss:.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)
+    with torch.no_grad():
+        logits, _ = model(x)
+        probs = torch.softmax(logits.squeeze() / temperature, dim=0).cpu().numpy()
+        return np.random.choice(vocab, p=probs)
+
+###################################
+# Train and evaluate
+###################################
+
+if train_button:
+    st.write(f"Training **{model_type}** model...")
+
+    if model_type == "N-gram":
+        with st.spinner("Training N-gram model..."):
+            model = NGramModel(tokens, n=3)
+    elif model_type == "Feed Forward NN":
+        model = train_ffnn(tokens)
+    elif model_type == "Decision Tree":
+        model = train_dt(tokens)
+    elif model_type == "Gradient Boosted Tree":
+        model = train_gbt(tokens)
+    elif model_type == "RNN":
+        model = train_rnn(tokens)
+
+    st.session_state["model"] = model
+    st.session_state["model_type"] = model_type
+    st.success(f"{model_type} model trained.")
+
+###################################
+# Chat interface
+###################################
+
+st.header("💬 Chat with the model")
+
+if "model" in st.session_state:
+    user_input = st.text_input("Type a prompt:")
+
+    if user_input:
+        context = tokenize(user_input, tokenizer_type)
+        generated = context.copy()
+
+        for _ in range(20):
+            if st.session_state["model_type"] == "N-gram":
+                next_tok = st.session_state["model"].predict(generated, temperature)
+            elif st.session_state["model_type"] == "Feed Forward NN":
+                next_tok = ffnn_predict(st.session_state["model"], generated, temperature)
+            elif st.session_state["model_type"] == "Decision Tree":
+                next_tok = dt_predict(st.session_state["model"], generated)
+            elif st.session_state["model_type"] == "Gradient Boosted Tree":
+                next_tok = gbt_predict(st.session_state["model"], generated)
+            elif st.session_state["model_type"] == "RNN":
+                next_tok = rnn_predict(st.session_state["model"], generated, temperature)
+
+            generated.append(next_tok)
+            if next_tok == "<END>":
+                break
+
+        if tokenizer_type == "character":
+            output = "".join(generated)
+        else:
+            output = " ".join(generated)
+
+        st.write("**Model Output:**")
+        st.write(output)
+else:
+    st.info("Train a model to begin chatting.")