Update app.py

app.py

@@ -4,12 +4,12 @@ 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.
+ - full back-and-forth memory between turns.
 
 Requirements:
- - transformers
- - gradio
- - torch
+    pip install torch transformers gradio
 """
+
 import threading
 import gradio as gr
 import torch
@@ -20,7 +20,7 @@ MODEL_ID = "EpistemeAI/metatune-gpt20b-R0"
 print("Loading tokenizer and model (this may take a while)...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# Use auto dtype & device mapping as requested
+# Use auto dtype & device mapping
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype="auto",
@@ -54,7 +54,6 @@ def generate_stream(prompt: str, max_tokens: int, temperature: float, top_p: flo
     Stream partial strings via TextIteratorStreamer.
     """
     inputs = tokenizer(prompt, return_tensors="pt")
-    # 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:
@@ -82,74 +81,76 @@ def generate_stream(prompt: str, max_tokens: int, temperature: float, top_p: flo
 
 def visible_messages_from_history(real_history: list, streaming_partial: str | None):
     """
-    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).
+    Convert internal history into Gradio-visible messages.
+    - Show user messages.
+    - Show assistant messages (partial or final).
+    - Hide system messages.
     """
     msgs = []
     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"] = streaming_partial
         else:
-            # 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):
+def respond_stream(user_message, system_message, max_tokens, temperature, top_p, history_state):
     """
-    Gradio streaming handler:
-    - Append real user message + assistant placeholder to GLOBAL_HISTORY
-    - Yield visible message lists as the assistant generates tokens
+    Gradio streaming handler with persistent memory.
     """
-    # Add the user message and an assistant placeholder into the real history
+    if history_state is None:
+        history_state = []
+
+    # Sync local and global histories (optional global memory)
+    with HISTORY_LOCK:
+        GLOBAL_HISTORY[:] = history_state
+
+    # Add the new user message and placeholder assistant
     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
+        GLOBAL_HISTORY.append({"role": "assistant", "content": ""})
         snapshot = list(GLOBAL_HISTORY)
 
-    # Immediately show user message and assistant placeholder ("thinking...")
+    # Show initial "thinking..." state
     initial_display = visible_messages_from_history(snapshot, streaming_partial="thinking...")
-    yield initial_display
+    yield initial_display, snapshot
 
-    # Build prompt using the real history but exclude the last assistant placeholder's empty content
+    # Build prompt excluding assistant placeholder
     with HISTORY_LOCK:
-        prompt_history = [h for h in GLOBAL_HISTORY[:-1]]  # all except the placeholder assistant
+        prompt_history = [h for h in GLOBAL_HISTORY[:-1]]
     prompt = build_prompt(system_message or "", prompt_history, user_message or "")
 
-    # Stream generation and update the last assistant entry
+    # Stream generation and update assistant output
     for partial in generate_stream(prompt, max_tokens, temperature, top_p):
         with HISTORY_LOCK:
-            # update global last assistant content
             if GLOBAL_HISTORY and GLOBAL_HISTORY[-1]["role"] == "assistant":
                 GLOBAL_HISTORY[-1]["content"] = partial
             snapshot = list(GLOBAL_HISTORY)
         display = visible_messages_from_history(snapshot, streaming_partial=partial)
-        yield display
+        yield display, snapshot
 
-    # Finalize: ensure assistant final content is shown
+    # Final display
     with HISTORY_LOCK:
         final_snapshot = list(GLOBAL_HISTORY)
     final_display = visible_messages_from_history(final_snapshot, streaming_partial=final_snapshot[-1].get("content", ""))
-    yield final_display
+    yield final_display, final_snapshot
+
+
+def reset_all():
+    with HISTORY_LOCK:
+        GLOBAL_HISTORY.clear()
+    return [], []
 
 
 # --- Gradio UI ---
@@ -157,12 +158,13 @@ with gr.Blocks() as demo:
     gr.Markdown(f"**Model:** {MODEL_ID} — (system messages hidden; user visible)")
 
     chatbot = gr.Chatbot(elem_id="chatbot", label="Chat", type="messages", height=560)
+    history_state = gr.State([])
 
     with gr.Row():
         with gr.Column(scale=4):
             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 (hidden from UI)")
+            system_input = gr.Textbox(value="You are a Vibe Coder assistant.", label="System message (hidden)")
             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)")
@@ -170,20 +172,18 @@ with gr.Blocks() as demo:
 
     send_btn.click(
         fn=respond_stream,
-        inputs=[user_input, system_input, max_tokens, temperature, top_p],
-        outputs=[chatbot],
+        inputs=[user_input, system_input, max_tokens, temperature, top_p, history_state],
+        outputs=[chatbot, history_state],
         queue=True,
     )
 
     clear_btn = gr.Button("Reset conversation")
-    def reset_all():
-        with HISTORY_LOCK:
-            GLOBAL_HISTORY.clear()
-        return []
-    clear_btn.click(fn=reset_all, inputs=None, outputs=[chatbot])
+    clear_btn.click(fn=reset_all, inputs=None, outputs=[chatbot, history_state])
 
-    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).")
+    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()