Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,484 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import time
+import random
+from datetime import datetime
+from collections import Counter
+import re
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# ── Page config ────────────────────────────────────────────────────────────────
+st.set_page_config(
+    page_title="NanoChat · LLM Playground",
+    page_icon="⚡",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+# ── Custom CSS ─────────────────────────────────────────────────────────────────
+st.markdown("""
+<style>
+@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=Syne:wght@400;600;800&display=swap');
+
+html, body, [class*="css"] {
+    font-family: 'Syne', sans-serif;
+}
+
+/* Dark industrial background */
+.stApp {
+    background: #0d0d0f;
+    color: #e8e4dc;
+}
+
+/* Sidebar */
+section[data-testid="stSidebar"] {
+    background: #111114 !important;
+    border-right: 1px solid #2a2a30;
+}
+
+/* Headers */
+h1, h2, h3 {
+    font-family: 'Syne', sans-serif !important;
+    font-weight: 800 !important;
+    letter-spacing: -0.03em;
+}
+
+/* Chat messages */
+.chat-msg {
+    padding: 14px 18px;
+    border-radius: 4px;
+    margin: 8px 0;
+    font-family: 'Space Mono', monospace;
+    font-size: 0.85rem;
+    line-height: 1.7;
+    border-left: 3px solid transparent;
+}
+.chat-msg.user {
+    background: #1a1a1f;
+    border-left-color: #f0c040;
+    color: #e8e4dc;
+}
+.chat-msg.assistant {
+    background: #141418;
+    border-left-color: #4af0a0;
+    color: #c8f0dc;
+}
+.chat-msg .role-label {
+    font-size: 0.65rem;
+    letter-spacing: 0.15em;
+    text-transform: uppercase;
+    opacity: 0.5;
+    margin-bottom: 6px;
+}
+
+/* Metric cards */
+.metric-card {
+    background: #111114;
+    border: 1px solid #2a2a30;
+    border-radius: 4px;
+    padding: 16px 20px;
+    margin: 6px 0;
+}
+.metric-value {
+    font-family: 'Space Mono', monospace;
+    font-size: 2rem;
+    font-weight: 700;
+    color: #f0c040;
+    line-height: 1;
+}
+.metric-label {
+    font-size: 0.7rem;
+    letter-spacing: 0.12em;
+    text-transform: uppercase;
+    opacity: 0.45;
+    margin-top: 4px;
+}
+
+/* Input box override */
+.stTextInput > div > div > input, .stTextArea textarea {
+    background: #111114 !important;
+    color: #e8e4dc !important;
+    border: 1px solid #2a2a30 !important;
+    font-family: 'Space Mono', monospace !important;
+    font-size: 0.85rem !important;
+}
+.stTextInput > div > div > input:focus, .stTextArea textarea:focus {
+    border-color: #f0c040 !important;
+    box-shadow: 0 0 0 2px rgba(240,192,64,0.15) !important;
+}
+
+/* Buttons */
+.stButton > button {
+    background: #f0c040 !important;
+    color: #0d0d0f !important;
+    font-family: 'Syne', sans-serif !important;
+    font-weight: 700 !important;
+    font-size: 0.8rem !important;
+    letter-spacing: 0.08em !important;
+    text-transform: uppercase !important;
+    border: none !important;
+    border-radius: 2px !important;
+    padding: 10px 24px !important;
+}
+.stButton > button:hover {
+    background: #ffd760 !important;
+}
+
+/* Selectbox */
+.stSelectbox > div > div {
+    background: #111114 !important;
+    border: 1px solid #2a2a30 !important;
+    color: #e8e4dc !important;
+}
+
+/* Tabs */
+.stTabs [data-baseweb="tab-list"] {
+    background: transparent;
+    border-bottom: 1px solid #2a2a30;
+    gap: 0;
+}
+.stTabs [data-baseweb="tab"] {
+    font-family: 'Syne', sans-serif !important;
+    font-weight: 600 !important;
+    font-size: 0.78rem !important;
+    letter-spacing: 0.1em !important;
+    text-transform: uppercase !important;
+    color: #888 !important;
+    background: transparent !important;
+    border: none !important;
+    padding: 10px 24px !important;
+}
+.stTabs [aria-selected="true"] {
+    color: #f0c040 !important;
+    border-bottom: 2px solid #f0c040 !important;
+}
+
+/* Divider */
+hr {
+    border-color: #2a2a30 !important;
+}
+
+/* Spinner text */
+.stSpinner > div {
+    color: #4af0a0 !important;
+}
+
+/* Scrollbar */
+::-webkit-scrollbar { width: 4px; }
+::-webkit-scrollbar-track { background: #0d0d0f; }
+::-webkit-scrollbar-thumb { background: #2a2a30; border-radius: 2px; }
+
+/* Word freq bars */
+.word-bar-container { margin: 4px 0; }
+.word-bar-label {
+    font-family: 'Space Mono', monospace;
+    font-size: 0.72rem;
+    color: #888;
+    display: flex;
+    justify-content: space-between;
+    margin-bottom: 2px;
+}
+.word-bar {
+    height: 6px;
+    background: linear-gradient(90deg, #4af0a0, #f0c040);
+    border-radius: 1px;
+}
+</style>
+""", unsafe_allow_html=True)
+
+# ── Session state ──────────────────────────────────────────────────────────────
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "model_loaded" not in st.session_state:
+    st.session_state.model_loaded = False
+if "pipeline" not in st.session_state:
+    st.session_state.pipeline = None
+if "total_tokens" not in st.session_state:
+    st.session_state.total_tokens = 0
+if "response_times" not in st.session_state:
+    st.session_state.response_times = []
+if "turn_count" not in st.session_state:
+    st.session_state.turn_count = 0
+
+# ── Model loader ───────────────────────────────────────────────────────────────
+@st.cache_resource(show_spinner=False)
+def load_model(model_id: str):
+    from transformers import pipeline as hf_pipeline
+    pipe = hf_pipeline(
+        "text-generation",
+        model=model_id,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    return pipe
+
+# ── Helpers ────────────────────────────────────────────────────────────────────
+MODEL_OPTIONS = {
+    "SmolLM2-135M-Instruct (HF)": "HuggingFaceTB/SmolLM2-135M-Instruct",
+    "SmolLM2-360M-Instruct (HF)": "HuggingFaceTB/SmolLM2-360M-Instruct",
+    "TinyLlama-1.1B-Chat": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "Qwen2.5-0.5B-Instruct": "Qwen/Qwen2.5-0.5B-Instruct",
+}
+
+def count_tokens_approx(text: str) -> int:
+    return max(1, len(text.split()) * 4 // 3)
+
+def get_word_freq(messages, top_n=10):
+    all_text = " ".join(m["content"] for m in messages).lower()
+    words = re.findall(r"\b[a-z]{4,}\b", all_text)
+    stopwords = {"that","this","with","from","have","will","been","they",
+                 "what","when","your","just","more","also","some","than",
+                 "then","there","their","these","those","about","which","would"}
+    words = [w for w in words if w not in stopwords]
+    return Counter(words).most_common(top_n)
+
+def format_chat_history(messages, model_id: str):
+    """Build a prompt string compatible with most instruct models."""
+    if "SmolLM2" in model_id or "Qwen" in model_id:
+        # ChatML format
+        prompt = ""
+        for m in messages:
+            role = m["role"]
+            content = m["content"]
+            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+        prompt += "<|im_start|>assistant\n"
+    else:
+        # TinyLlama / Llama-2 chat format
+        prompt = "<s>"
+        for m in messages:
+            if m["role"] == "user":
+                prompt += f"[INST] {m['content']} [/INST]"
+            else:
+                prompt += f" {m['content']} </s><s>"
+    return prompt
+
+def generate_response(pipe, messages, model_id, max_new_tokens, temperature):
+    prompt = format_chat_history(messages, model_id)
+    t0 = time.time()
+    out = pipe(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        do_sample=temperature > 0,
+        pad_token_id=pipe.tokenizer.eos_token_id,
+        return_full_text=False,
+    )
+    elapsed = time.time() - t0
+    text = out[0]["generated_text"].strip()
+    # Strip any trailing special tokens
+    for tok in ["<|im_end|>", "</s>", "[INST]"]:
+        text = text.split(tok)[0].strip()
+    return text, elapsed
+
+# ── Sidebar ────────────────────────────────────────────────────────────────────
+with st.sidebar:
+    st.markdown("## ⚡ NanoChat")
+    st.markdown("<p style='font-size:0.75rem;color:#666;font-family:Space Mono,monospace;margin-top:-8px'>Open-weight LLM Playground</p>", unsafe_allow_html=True)
+    st.divider()
+
+    selected_label = st.selectbox("Model", list(MODEL_OPTIONS.keys()))
+    model_id = MODEL_OPTIONS[selected_label]
+
+    max_new_tokens = st.slider("Max new tokens", 32, 512, 200, 16)
+    temperature = st.slider("Temperature", 0.0, 1.5, 0.7, 0.05)
+
+    st.divider()
+
+    if st.button("⚡  Load / Reload Model"):
+        with st.spinner(f"Loading {selected_label}…"):
+            try:
+                st.session_state.pipeline = load_model(model_id)
+                st.session_state.model_loaded = True
+                st.success("Model ready!")
+            except Exception as e:
+                st.error(f"Error: {e}")
+
+    if st.button("🗑  Clear Chat"):
+        st.session_state.messages = []
+        st.session_state.total_tokens = 0
+        st.session_state.response_times = []
+        st.session_state.turn_count = 0
+        st.rerun()
+
+    st.divider()
+    st.markdown(f"""
+    <div style='font-family:Space Mono,monospace;font-size:0.68rem;color:#555;line-height:2'>
+    Model ID<br>
+    <span style='color:#f0c040'>{model_id.split("/")[-1]}</span><br><br>
+    Status<br>
+    <span style='color:{"#4af0a0" if st.session_state.model_loaded else "#f06060"}'>
+    {"● Loaded" if st.session_state.model_loaded else "○ Not loaded"}
+    </span>
+    </div>
+    """, unsafe_allow_html=True)
+
+# ── Main area ──────────────────────────────────────────────────────────────────
+tab_chat, tab_viz = st.tabs(["💬  Chat", "📊  Analytics"])
+
+# ─── Chat tab ──────────────────────────────────────────────────────────────────
+with tab_chat:
+    st.markdown("<h1 style='margin-bottom:2px'>Chat</h1>", unsafe_allow_html=True)
+    st.markdown(f"<p style='font-size:0.78rem;color:#555;font-family:Space Mono,monospace;margin-bottom:24px'>{model_id}</p>", unsafe_allow_html=True)
+
+    if not st.session_state.model_loaded:
+        st.info("👈  Load a model from the sidebar to begin.")
+    else:
+        # Render history
+        chat_container = st.container()
+        with chat_container:
+            for msg in st.session_state.messages:
+                role_label = "YOU" if msg["role"] == "user" else "AI"
+                css_class = "user" if msg["role"] == "user" else "assistant"
+                st.markdown(f"""
+                <div class='chat-msg {css_class}'>
+                  <div class='role-label'>{role_label}</div>
+                  {msg['content']}
+                </div>
+                """, unsafe_allow_html=True)
+
+        # Input
+        with st.form("chat_form", clear_on_submit=True):
+            cols = st.columns([8, 1])
+            with cols[0]:
+                user_input = st.text_area("Message", height=80, label_visibility="collapsed",
+                                          placeholder="Type a message and press Send…")
+            with cols[1]:
+                submitted = st.form_submit_button("Send", use_container_width=True)
+
+        if submitted and user_input.strip():
+            st.session_state.messages.append({"role": "user", "content": user_input.strip()})
+            st.session_state.total_tokens += count_tokens_approx(user_input)
+
+            with st.spinner("Thinking…"):
+                try:
+                    reply, elapsed = generate_response(
+                        st.session_state.pipeline,
+                        st.session_state.messages,
+                        model_id,
+                        max_new_tokens,
+                        temperature,
+                    )
+                    st.session_state.messages.append({"role": "assistant", "content": reply})
+                    st.session_state.total_tokens += count_tokens_approx(reply)
+                    st.session_state.response_times.append(round(elapsed, 2))
+                    st.session_state.turn_count += 1
+                except Exception as e:
+                    st.error(f"Generation error: {e}")
+            st.rerun()
+
+# ─── Analytics tab ─────────────────────────────────────────────────────────────
+with tab_viz:
+    st.markdown("<h1 style='margin-bottom:2px'>Analytics</h1>", unsafe_allow_html=True)
+    st.markdown("<p style='font-size:0.78rem;color:#555;font-family:Space Mono,monospace;margin-bottom:24px'>Session insights</p>", unsafe_allow_html=True)
+
+    msgs = st.session_state.messages
+    rt   = st.session_state.response_times
+
+    # ── Metrics row ────────────────────────────────────────────────────────────
+    c1, c2, c3, c4 = st.columns(4)
+    with c1:
+        st.markdown(f"""
+        <div class='metric-card'>
+          <div class='metric-value'>{st.session_state.turn_count}</div>
+          <div class='metric-label'>Turns</div>
+        </div>""", unsafe_allow_html=True)
+    with c2:
+        st.markdown(f"""
+        <div class='metric-card'>
+          <div class='metric-value'>{st.session_state.total_tokens}</div>
+          <div class='metric-label'>Est. Tokens</div>
+        </div>""", unsafe_allow_html=True)
+    with c3:
+        avg_rt = round(sum(rt)/len(rt), 2) if rt else 0.0
+        st.markdown(f"""
+        <div class='metric-card'>
+          <div class='metric-value'>{avg_rt}s</div>
+          <div class='metric-label'>Avg Response</div>
+        </div>""", unsafe_allow_html=True)
+    with c4:
+        user_msgs = [m for m in msgs if m["role"]=="user"]
+        avg_len = round(sum(len(m["content"].split()) for m in user_msgs)/len(user_msgs)) if user_msgs else 0
+        st.markdown(f"""
+        <div class='metric-card'>
+          <div class='metric-value'>{avg_len}</div>
+          <div class='metric-label'>Avg User Words</div>
+        </div>""", unsafe_allow_html=True)
+
+    st.divider()
+
+    col_left, col_right = st.columns(2)
+
+    # ── Response time chart ────────────────────────────────────────────────────
+    with col_left:
+        st.markdown("#### Response times (s)")
+        if rt:
+            import pandas as pd
+            df_rt = pd.DataFrame({"Turn": list(range(1, len(rt)+1)), "Seconds": rt})
+            st.line_chart(df_rt.set_index("Turn"), color="#4af0a0", height=200)
+        else:
+            st.caption("No data yet — start chatting!")
+
+    # ── Message length chart ───────────────────────────────────────────────────
+    with col_right:
+        st.markdown("#### Message lengths (words)")
+        if msgs:
+            import pandas as pd
+            rows = []
+            u_idx = a_idx = 1
+            for m in msgs:
+                wc = len(m["content"].split())
+                if m["role"] == "user":
+                    rows.append({"idx": u_idx, "role": "User", "words": wc})
+                    u_idx += 1
+                else:
+                    rows.append({"idx": a_idx, "role": "AI", "words": wc})
+                    a_idx += 1
+            import pandas as pd
+            df_ml = pd.DataFrame(rows)
+            st.bar_chart(df_ml.pivot_table(index="idx", columns="role", values="words", aggfunc="sum").fillna(0),
+                         color=["#f0c040", "#4af0a0"], height=200)
+        else:
+            st.caption("No data yet — start chatting!")
+
+    st.divider()
+
+    # ── Word frequency ──────────────────────────────────────────────────────────
+    st.markdown("#### Top words across conversation")
+    if msgs:
+        freq = get_word_freq(msgs, top_n=12)
+        if freq:
+            max_count = freq[0][1]
+            for word, count in freq:
+                pct = int((count / max_count) * 100)
+                st.markdown(f"""
+                <div class='word-bar-container'>
+                  <div class='word-bar-label'><span>{word}</span><span>{count}</span></div>
+                  <div class='word-bar' style='width:{pct}%'></div>
+                </div>""", unsafe_allow_html=True)
+    else:
+        st.caption("No data yet — start chatting!")
+
+    st.divider()
+
+    # ── Role distribution ────────────────────────────────────────────────────
+    st.markdown("#### Message distribution")
+    if msgs:
+        u_count = sum(1 for m in msgs if m["role"]=="user")
+        a_count = sum(1 for m in msgs if m["role"]=="assistant")
+        total   = u_count + a_count
+        u_pct   = int(u_count/total*100)
+        a_pct   = 100 - u_pct
+        st.markdown(f"""
+        <div style='display:flex;gap:0;border-radius:3px;overflow:hidden;height:28px;margin:8px 0'>
+          <div style='width:{u_pct}%;background:#f0c040;display:flex;align-items:center;
+                      justify-content:center;font-family:Space Mono,monospace;
+                      font-size:0.7rem;color:#0d0d0f;font-weight:700'>
+            USER {u_pct}%
+          </div>
+          <div style='width:{a_pct}%;background:#4af0a0;display:flex;align-items:center;
+                      justify-content:center;font-family:Space Mono,monospace;
+                      font-size:0.7rem;color:#0d0d0f;font-weight:700'>
+            AI {a_pct}%
+          </div>
+        </div>
+        """, unsafe_allow_html=True)
+    else:
+        st.caption("No data yet.")