Update app.py

app.py

@@ -1,29 +1,46 @@
 # 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.
+
+Requirements:
+ - transformers
+ - accelerate (recommended)
+ - gradio
+ - torch
+"""
 import threading
+import time
 import gradio as gr
 import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    TextIteratorStreamer,
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
 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 mapping / dtype
+# ALWAYS use auto for dtype & device_map as requested
+# This will let HF/accelerate place weights across available devices
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype="auto",
     device_map="auto",
 )
 model.eval()
-print("Model loaded. Device:", next(model.parameters()).device)
+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 = []
+HISTORY_LOCK = threading.Lock()
+
 
-# --------- Helper: build prompt ----------
-def build_prompt(system_message: str, history: list[dict], user_message: str) -> str:
+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.
+    """
     pieces = []
     if system_message:
         pieces.append(f"<|system|>\n{system_message}\n")
@@ -34,10 +51,19 @@ def build_prompt(system_message: str, history: list[dict], user_message: str) ->
     pieces.append(f"<|user|>\n{user_message}\n<|assistant|>\n")
     return "\n".join(pieces)
 
-# --------- Streaming generator ----------
-def generate_stream(prompt, max_tokens, temperature, top_p):
+
+def generate_stream(prompt: str, max_tokens: int, temperature: float, top_p: float):
+    """
+    Yields partial strings as the model generates tokens using TextIteratorStreamer.
+    """
+    # Tokenize (we avoid forcing a single-device .to(...) in case of HF sharded device_map)
     inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"].to(model.device)
+    # Move input_ids to same device as a model parameter (works with many configs)
+    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)
 
@@ -58,56 +84,119 @@ def generate_stream(prompt, max_tokens, temperature, top_p):
         partial += token_str
         yield partial
 
-# --------- Gradio app logic ----------
-def respond_stream(user_message, chat_history, system_message, max_tokens, temperature, top_p):
-    history = chat_history or []
-    prompt = build_prompt(system_message or "", history, user_message or "")
-
-    history.append({"role": "user", "content": user_message})
-    history.append({"role": "assistant", "content": ""})
 
-    def history_to_chatbot_rows(hist):
-        rows = []
-        for item in hist:
-            if item["role"] == "assistant":
-                rows.append(("thinking...", item["content"] or "thinking..."))
-        return rows or []
+def make_display_messages_from_history(real_history: list, current_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.
+    """
+    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
+        if msgs and msgs[-1]["role"] == "assistant":
+            msgs[-1]["content"] = current_partial
+        else:
+            msgs.append({"role": "user", "content": "thinking..."})
+            msgs.append({"role": "assistant", "content": current_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.
+    """
+    # Append the real user turn and an empty assistant placeholder to GLOBAL_HISTORY
+    with HISTORY_LOCK:
+        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)
+
+    # initial UI placeholder: show existing assistant rows and the new placeholder
+    displayed = make_display_messages_from_history(local_history, current_partial="thinking...")
+    yield displayed
+
+    # 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)
+    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 = build_prompt(system_message or "", prompt_history, user_message or "")
+
+    # Stream generation
+    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
+            if GLOBAL_HISTORY and GLOBAL_HISTORY[-1]["role"] == "assistant":
+                GLOBAL_HISTORY[-1]["content"] = partial
+            current_snapshot = list(GLOBAL_HISTORY)
 
-    chatbot_rows = history_to_chatbot_rows(history[:-1])
-    chatbot_rows.append(("thinking...", "thinking..."))
-    yield chatbot_rows  # placeholder row
+        displayed = make_display_messages_from_history(current_snapshot, current_partial=partial)
+        yield displayed
 
-    for partial in generate_stream(prompt, max_tokens, temperature, top_p):
-        chatbot_rows[-1] = ("thinking...", partial)
-        history[-1]["content"] = partial
-        yield chatbot_rows
+    # final sync: ensure the assistant content is finalized in GLOBAL_HISTORY (already done)
+    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
 
-    chatbot_rows[-1] = ("thinking...", history[-1]["content"])
-    yield chatbot_rows
 
-# --------- Build Gradio UI ----------
+# --- Gradio UI ---
 with gr.Blocks() as demo:
-    gr.Markdown(f"**Model:** {MODEL_ID}")
-    with gr.Row():
-        chatbot = gr.Chatbot(elem_id="chatbot", label="Assistant Output (user/system hidden)").style(height=500)
+    gr.Markdown(f"**Model:** {MODEL_ID} — (UI hides user/system messages; left column shows 'thinking...')")
 
-    history_state = gr.State(value=[])
-    system_input = gr.Textbox(value="You are a Vibe Coder assistant.", label="System message")
-    user_input = gr.Textbox(placeholder="Type a user message and press Send", label="Your message")
-    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")
+    # 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)
 
+    with gr.Row():
+        with gr.Column(scale=4):
+            user_input = gr.Textbox(placeholder="Type a user 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")
+            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, history_state, system_input, max_tokens, temperature, top_p],
+        inputs=[user_input, system_input, max_tokens, temperature, top_p],
         outputs=[chatbot],
         queue=True,
     )
 
-    send_btn.click(lambda u, s: s, inputs=[user_input, history_state], outputs=[history_state])
-    send_btn.click(lambda: "", None, user_input)
+    # Optional controls
+    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])
+
+    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.")
 
 if __name__ == "__main__":
     demo.launch()