import streamlit as st
import time
import random
from datetime import datetime
from collections import Counter
import re
# ── Page config ────────────────────────────────────────────────────────────────
st.set_page_config(
page_title="NanoChat · LLM Playground",
page_icon="⚡",
layout="wide",
initial_sidebar_state="expanded",
)
# ── Custom CSS ─────────────────────────────────────────────────────────────────
st.markdown("""
""", 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 = ""
for m in messages:
if m["role"] == "user":
prompt += f"[INST] {m['content']} [/INST]"
else:
prompt += f" {m['content']} "
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|>", "", "[INST]"]:
text = text.split(tok)[0].strip()
return text, elapsed
# ── Sidebar ────────────────────────────────────────────────────────────────────
with st.sidebar:
st.markdown("## ⚡ NanoChat")
st.markdown("Open-weight LLM Playground
", 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"""
Model ID
{model_id.split("/")[-1]}
Status
{"● Loaded" if st.session_state.model_loaded else "○ Not loaded"}
""", unsafe_allow_html=True)
# ── Main area ──────────────────────────────────────────────────────────────────
tab_chat, tab_viz = st.tabs(["💬 Chat", "📊 Analytics"])
# ─── Chat tab ──────────────────────────────────────────────────────────────────
with tab_chat:
st.markdown("Chat
", unsafe_allow_html=True)
st.markdown(f"{model_id}
", 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"""
{role_label}
{msg['content']}
""", 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("Analytics
", unsafe_allow_html=True)
st.markdown("Session insights
", 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"""
{st.session_state.turn_count}
Turns
""", unsafe_allow_html=True)
with c2:
st.markdown(f"""
{st.session_state.total_tokens}
Est. Tokens
""", unsafe_allow_html=True)
with c3:
avg_rt = round(sum(rt)/len(rt), 2) if rt else 0.0
st.markdown(f"""
""", 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"""
""", 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"""
""", 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"""
USER {u_pct}%
AI {a_pct}%
""", unsafe_allow_html=True)
else:
st.caption("No data yet.")