Deploy transformer demo to Space

app.py

@@ -1,18 +1,19 @@
 import streamlit as st
 from assist.chat import chat as embed_chat
 from assist.bayes_chat import bayes_chat
+from assist.transformer_demo import transformer_next
 
 st.title("RepoSage Chatbot Demo")
 
 question = st.text_input("Enter your question below:")
 
-col1, col2 = st.columns(2)
+col1, col2, col3 = st.columns(3)
 with col1:
-    if st.button("Ask Embedding RepoSage"):
-        answer = embed_chat(question)
-        st.write(answer)
-
+    if st.button("Embedding Q&A"):
+        st.write(embed_chat(question))
 with col2:
-    if st.button("Ask Bayesian RepoSage"):
-        answer = bayes_chat(question)
-        st.write(answer)
+    if st.button("Bayesian Q&A"):
+        st.write(bayes_chat(question))
+with col3:
+    if st.button("Transformer Demo"):
+        st.write(transformer_next(question))

assist/main.py

@@ -22,5 +22,11 @@ def chat(question: str = typer.Argument(..., help="Question to ask RepoSage")):
     response = chat_plugin(question)
     print(response)
 
+@app.command()
+def transform(prompt: str = typer.Argument(..., help="Prompt for transformer demo")):
+    """Invoke the single‐block transformer next-token demo."""
+    from .transformer_demo import transformer_next
+    print(transformer_next(prompt))
+
 if __name__ == "__main__":
     app()

assist/transformer_demo.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from .chat import TOKEN2IDX, IDX2TOKEN  # reuse your vocab maps
+from .chat import WEIGHTS              # reuse your embedding weights
+
+class SingleTransformerBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads=2):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.ff    = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * 4),
+            nn.ReLU(),
+            nn.Linear(embed_dim * 4, embed_dim)
+        )
+        self.norm2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, x):
+        # Self-attention
+        attn_out, _ = self.attn(x, x, x)
+        x = self.norm1(x + attn_out)
+        # Feed-forward
+        ff_out = self.ff(x)
+        x = self.norm2(x + ff_out)
+        return x
+
+# Instantiate once
+_EMB = torch.tensor(WEIGHTS, dtype=torch.float32)              # V×D
+_block = SingleTransformerBlock(embed_dim=_EMB.size(1), num_heads=2)
+
+def transformer_next(prompt: str) -> str:
+    """
+    Given a prompt, tokenize it, embed each token, run through one
+    transformer block, then use the last position’s output vector
+    to pick the nearest vocab token as the “next token.”
+    """
+    tokens = prompt.lower().split()
+    idxs = [TOKEN2IDX[t] for t in tokens if t in TOKEN2IDX]
+    if not idxs:
+        return "🤔 No known tokens to predict from."
+    # Build batch: 1×seq_len×D
+    x = _EMB[idxs].unsqueeze(0)
+    # Forward pass
+    out = _block(x)              # 1×seq_len×D
+    last = out[0, -1].unsqueeze(0)  # 1×D
+    # Cosine similarity against all embeddings
+    sims = nn.functional.cosine_similarity(last, _EMB)
+    best = int(torch.argmax(sims))
+    return f"🔮 Next‐token prediction: **{IDX2TOKEN[best]}**"