Update app.py

app.py

@@ -1,17 +1,16 @@
 # save as app.py
 """
-Gradio streaming chat that hides user/system text.
-Left-side messages are always "thinking..." (literal).
-Right-side shows the assistant output streamed as it is generated.
+Gradio streaming chat where:
+ - user messages are visible in the UI,
+ - system messages are hidden (kept for context),
+ - assistant output is streamed and updates in-place.
 
 Requirements:
  - transformers
- - accelerate (recommended)
  - gradio
  - torch
 """
 import threading
-import time
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
@@ -21,8 +20,7 @@ MODEL_ID = "EpistemeAI/gpt-oss-20b-RL"
 print("Loading tokenizer and model (this may take a while)...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# ALWAYS use auto for dtype & device_map as requested
-# This will let HF/accelerate place weights across available devices
+# Use auto dtype & device mapping as requested
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype="auto",
@@ -31,15 +29,14 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 print("Model loaded. Example param device:", next(model.parameters()).device)
 
-# Global history (real content). Stored as list of {"role": "user"|"assistant"|"system", "content": "..."}
-GLOBAL_HISTORY = []
+# Thread-safe global history
+GLOBAL_HISTORY = []  # list of {"role": "system"|"user"|"assistant", "content": "..."}
 HISTORY_LOCK = threading.Lock()
 
 
 def build_prompt(system_message: str, history: list, user_message: str) -> str:
     """
-    Build the model prompt in your preferred format.
-    Adjust this function if your model expects a different conversation format.
+    Build prompt in the model's expected format. Adjust as needed.
     """
     pieces = []
     if system_message:
@@ -54,15 +51,13 @@ def build_prompt(system_message: str, history: list, user_message: str) -> str:
 
 def generate_stream(prompt: str, max_tokens: int, temperature: float, top_p: float):
     """
-    Yields partial strings as the model generates tokens using TextIteratorStreamer.
+    Stream partial strings via TextIteratorStreamer.
     """
-    # Tokenize (we avoid forcing a single-device .to(...) in case of HF sharded device_map)
     inputs = tokenizer(prompt, return_tensors="pt")
-    # Move input_ids to same device as a model parameter (works with many configs)
+    # Move input ids to model param device where possible (works with many accelerate setups)
     try:
         input_ids = inputs["input_ids"].to(next(model.parameters()).device)
     except Exception:
-        # Fallback: do not move if that fails (accelerate may handle placement)
         input_ids = inputs["input_ids"]
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
@@ -85,101 +80,94 @@ def generate_stream(prompt: str, max_tokens: int, temperature: float, top_p: flo
         yield partial
 
 
-def make_display_messages_from_history(real_history: list, current_partial: str | None):
+def visible_messages_from_history(real_history: list, streaming_partial: str | None):
     """
-    Convert the internal 'real_history' (which contains real user & assistant content)
-    into the list of openai-style message dicts that Gradio Chatbot (type="messages")
-    expects. Every non-assistant message is replaced with a literal "thinking...".
-    For each assistant exchange we produce:
-      {"role":"user", "content":"thinking..."}
-      {"role":"assistant", "content": "<assistant content or partial stream>"}
-    The UI will therefore show the left side text "thinking..." and right side the assistant.
+    Convert internal history into the list of OpenAI-style messages for Gradio UI.
+
+    - Show user messages verbatim (visible).
+    - Show assistant messages (streamed or final).
+    - Omit system messages (kept only for model context).
     """
     msgs = []
-    # Walk the real history and whenever we hit an assistant turn, pair it with a thinking user
-    i = 0
-    while i < len(real_history):
-        item = real_history[i]
-        if item["role"] == "assistant":
-            # the user message that caused this assistant reply is typically just before it,
-            # but we hide the real user content and show "thinking..." instead.
-            msgs.append({"role": "user", "content": "thinking..."})
-            msgs.append({"role": "assistant", "content": item.get("content", "") or "thinking..."})
-        i += 1
-
-    # If a current_partial exists (we're streaming a new assistant response),
-    # ensure it's reflected as the last assistant message (with a preceding "thinking...")
-    if current_partial is not None:
-        # If the last two entries are already the streaming pair, replace them; otherwise append new
+    for entry in real_history:
+        role = entry.get("role")
+        content = entry.get("content", "")
+        if role == "system":
+            # hide system from UI
+            continue
+        # For assistant messages, we'll use content (may be empty)
+        msgs.append({"role": role, "content": content or ("thinking..." if role == "assistant" else "")})
+
+    # If we're currently streaming an assistant response, ensure it's reflected as the last assistant msg
+    if streaming_partial is not None:
+        # If last message is assistant, replace its content, otherwise append a new (user, assistant) pair
         if msgs and msgs[-1]["role"] == "assistant":
-            msgs[-1]["content"] = current_partial
+            msgs[-1]["content"] = streaming_partial
         else:
-            msgs.append({"role": "user", "content": "thinking..."})
-            msgs.append({"role": "assistant", "content": current_partial})
+            # The user message that started this assistant reply should already be in history and visible.
+            # Append assistant partial as the reply
+            msgs.append({"role": "assistant", "content": streaming_partial})
+
     return msgs
 
 
 def respond_stream(user_message, system_message, max_tokens, temperature, top_p):
     """
-    Gradio streaming function that yields successive message-lists (OpenAI-style dicts).
-    It mutates GLOBAL_HISTORY to store the true conversation, but the UI only ever sees
-    'thinking...' in non-assistant slots and the assistant's streamed content on the right.
+    Gradio streaming handler:
+    - Append real user message + assistant placeholder to GLOBAL_HISTORY
+    - Yield visible message lists as the assistant generates tokens
     """
-    # Append the real user turn and an empty assistant placeholder to GLOBAL_HISTORY
+    # Add the user message and an assistant placeholder into the real history
     with HISTORY_LOCK:
+        if system_message:
+            # include system message in real history for model context (but it won't be shown)
+            GLOBAL_HISTORY.append({"role": "system", "content": system_message})
         GLOBAL_HISTORY.append({"role": "user", "content": user_message})
-        GLOBAL_HISTORY.append({"role": "assistant", "content": ""})  # placeholder for streaming
-        # create a shallow copy for local read
-        local_history = list(GLOBAL_HISTORY)
+        GLOBAL_HISTORY.append({"role": "assistant", "content": ""})  # placeholder
+        snapshot = list(GLOBAL_HISTORY)
 
-    # initial UI placeholder: show existing assistant rows and the new placeholder
-    displayed = make_display_messages_from_history(local_history, current_partial="thinking...")
-    yield displayed
+    # Immediately show user message and assistant placeholder ("thinking...")
+    initial_display = visible_messages_from_history(snapshot, streaming_partial="thinking...")
+    yield initial_display
 
-    # Build model prompt from the real history (exclude the last assistant placeholder content)
-    # We pass the actual global history (safe to read under lock copy)
+    # Build prompt using the real history but exclude the last assistant placeholder's empty content
     with HISTORY_LOCK:
-        # Send a snapshot (exclude the last assistant placeholder since it's empty)
-        prompt_history = [h for h in GLOBAL_HISTORY[:-1] if h.get("role")]
+        prompt_history = [h for h in GLOBAL_HISTORY[:-1]]  # all except the placeholder assistant
     prompt = build_prompt(system_message or "", prompt_history, user_message or "")
 
-    # Stream generation
+    # Stream generation and update the last assistant entry
     for partial in generate_stream(prompt, max_tokens, temperature, top_p):
-        # Update the global assistant placeholder with the partial so future turns keep context
         with HISTORY_LOCK:
-            # Update the last assistant placeholder
+            # update global last assistant content
             if GLOBAL_HISTORY and GLOBAL_HISTORY[-1]["role"] == "assistant":
                 GLOBAL_HISTORY[-1]["content"] = partial
-            current_snapshot = list(GLOBAL_HISTORY)
-
-        displayed = make_display_messages_from_history(current_snapshot, current_partial=partial)
-        yield displayed
+            snapshot = list(GLOBAL_HISTORY)
+        display = visible_messages_from_history(snapshot, streaming_partial=partial)
+        yield display
 
-    # final sync: ensure the assistant content is finalized in GLOBAL_HISTORY (already done)
+    # Finalize: ensure assistant final content is shown
     with HISTORY_LOCK:
         final_snapshot = list(GLOBAL_HISTORY)
-    displayed = make_display_messages_from_history(final_snapshot, current_partial=final_snapshot[-1].get("content", ""))
-    yield displayed
+    final_display = visible_messages_from_history(final_snapshot, streaming_partial=final_snapshot[-1].get("content", ""))
+    yield final_display
 
 
 # --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown(f"**Model:** {MODEL_ID} — (UI hides user/system messages; left column shows 'thinking...')")
+    gr.Markdown(f"**Model:** {MODEL_ID} — (system messages hidden; user visible)")
 
-    # Chatbot expects a list of {"role":.., "content":..} dicts when type="messages"
-    chatbot = gr.Chatbot(elem_id="chatbot", label="Assistant (user/system hidden)", type="messages", height=560)
+    chatbot = gr.Chatbot(elem_id="chatbot", label="Chat", type="messages", height=560)
 
     with gr.Row():
         with gr.Column(scale=4):
-            user_input = gr.Textbox(placeholder="Type a user message and press Send", label="Your message")
+            user_input = gr.Textbox(placeholder="Type a message and press Send", label="Your message")
         with gr.Column(scale=2):
-            system_input = gr.Textbox(value="You are a Vibe Coder assistant.", label="System message")
+            system_input = gr.Textbox(value="You are a Vibe Coder assistant.", label="System message (hidden from UI)")
             max_tokens = gr.Slider(minimum=1, maximum=4000, value=800, step=1, label="Max new tokens")
             temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.01, label="Temperature")
             top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)")
             send_btn = gr.Button("Send")
 
-    # Hook the streaming respond function. Gradio will accept a generator that yields message lists.
     send_btn.click(
         fn=respond_stream,
         inputs=[user_input, system_input, max_tokens, temperature, top_p],
@@ -187,7 +175,6 @@ with gr.Blocks() as demo:
         queue=True,
     )
 
-    # Optional controls
     clear_btn = gr.Button("Reset conversation")
     def reset_all():
         with HISTORY_LOCK:
@@ -195,8 +182,8 @@ with gr.Blocks() as demo:
         return []
     clear_btn.click(fn=reset_all, inputs=None, outputs=[chatbot])
 
-    gr.Markdown("Note: model loading uses `device_map='auto'` and `torch_dtype='auto'`. "
-                "If you run into out-of-memory problems on small GPUs, consider running on a machine with more memory or using model parallel tools.")
+    gr.Markdown("Notes: model loading uses `device_map='auto'` and `torch_dtype='auto'`. "
+                "If running multi-worker (gunicorn) you will need an external history store (Redis/DB).")
 
 if __name__ == "__main__":
     demo.launch()