Add decoding controls, few-shot prompt, repetition guard

app.py

@@ -18,17 +18,36 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 
-def respond(message, history):
-    # Simpler prompt (byte tokenizer is sensitive to unfamiliar special tokens)
-    turns = []
+SYS_PROMPT = (
+    "You are a helpful assistant. Answer briefly and clearly. "
+    "Avoid repeating characters. If unsure, say 'I don't know'."
+)
+
+FEW_SHOTS = [
+    ("Hello", "Hello!"),
+    ("What is the capital of France?", "Paris."),
+    ("2+2?", "4."),
+]
+
+
+def build_prompt(message: str, history: list[tuple[str, str]]) -> str:
+    # Minimal, byte-tokenizer-friendly prompt (no special tokens)
+    lines = [f"System: {SYS_PROMPT}"]
+    for q, a in FEW_SHOTS:
+        lines.append(f"User: {q}")
+        lines.append(f"Assistant: {a}")
     for u, a in history:
         if u:
-            turns.append(f"User: {u}")
+            lines.append(f"User: {u}")
         if a:
-            turns.append(f"Assistant: {a}")
-    turns.append(f"User: {message}")
-    turns.append("Assistant:")
-    prompt = "\n".join(turns)
+            lines.append(f"Assistant: {a}")
+    lines.append(f"User: {message}")
+    lines.append("Assistant:")
+    return "\n".join(lines)
+
+
+def respond(message, history, max_new_tokens, temperature, top_p, top_k, repetition_penalty, no_repeat_ngram_size):
+    prompt = build_prompt(message, history)
     x = tok(prompt, return_tensors="pt")
     if torch.cuda.is_available():
         x = {k: v.to(model.device) for k, v in x.items()}
@@ -36,12 +55,13 @@ def respond(message, history):
     streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True)
     gen_kwargs = dict(
         **x,
-        max_new_tokens=128,
-        do_sample=True,
-        top_p=0.9,
-        temperature=0.7,
-        repetition_penalty=1.2,
-        no_repeat_ngram_size=3,
+        max_new_tokens=int(max_new_tokens),
+        do_sample=bool(temperature > 0.0),
+        top_p=float(top_p),
+        top_k=int(top_k),
+        temperature=max(1e-6, float(temperature)),
+        repetition_penalty=max(1.0, float(repetition_penalty)),
+        no_repeat_ngram_size=int(no_repeat_ngram_size),
         streamer=streamer,
     )
     thread = Thread(target=model.generate, kwargs=gen_kwargs)
@@ -49,17 +69,53 @@ def respond(message, history):
     partial = ""
     for new_text in streamer:
         partial += new_text
+        # Simple repetition guard: if too many identical trailing chars, stop early
+        tail = partial[-200:]
+        if len(tail) >= 10 and any(tail.endswith(c * 10) for c in set(tail)):
+            break
         yield partial
 
 
-demo = gr.ChatInterface(
-    fn=respond,
-    title="NSA 117M Chat (byte tokenizer)",
-    description=(
-        "Byte-level tokenizer (vocab=256). No KV cache in v1; streaming enabled."
-    ),
-    examples=[["Write a haiku about sparse attention."], ["Explain NSA branches succinctly."]],
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# NSA 117M Chat (byte tokenizer)")
+    gr.Markdown("Byte-level tokenizer (vocab=256). Streaming enabled. Use controls to reduce repetition.")
+    chat = gr.Chatbot()
+    with gr.Row():
+        msg = gr.Textbox(label="Message")
+    with gr.Accordion("Decoding controls", open=False):
+        max_new = gr.Slider(16, 512, value=128, step=16, label="Max new tokens")
+        temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature (0 = greedy)")
+        top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
+        top_k = gr.Slider(0, 200, value=50, step=10, label="Top-k (0 disables)")
+        rep_pen = gr.Slider(1.0, 2.0, value=1.3, step=0.05, label="Repetition penalty")
+        ngram = gr.Slider(0, 6, value=3, step=1, label="No-repeat n-gram size (0 disables)")
+
+    def user_submit(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot_respond(history, max_new_tokens, temperature, top_p, top_k, repetition_penalty, no_repeat_ngram_size):
+        user_message = history[-1][0]
+        gen = respond(
+            user_message,
+            [(u, a) for u, a in history[:-1] if u is not None and a is not None],
+            max_new_tokens,
+            temperature,
+            top_p,
+            top_k,
+            repetition_penalty,
+            no_repeat_ngram_size,
+        )
+        partial = ""
+        for part in gen:
+            partial = part
+            history[-1][1] = partial
+            yield history
+
+    msg.submit(user_submit, [msg, chat], [msg, chat]).then(
+        bot_respond,
+        [chat, max_new, temperature, top_p, top_k, rep_pen, ngram],
+        [chat],
+    )
 
 
 if __name__ == "__main__":