Add official TwIL-LM3 Gradio ZeroGPU demo

README.md

@@ -1,13 +1,34 @@
 ---
-title: TwIL LM3
-emoji: ⚡
-colorFrom: purple
-colorTo: pink
+title: TwIL-LM3
+emoji: 🧠
+colorFrom: indigo
+colorTo: blue
 sdk: gradio
-sdk_version: 6.24.0
-python_version: '3.12'
+sdk_version: 6.22.0
 app_file: app.py
-pinned: false
+pinned: true
+license: other
+short_description: Official demo of TwIL-LM3, a 3B formal-logic reasoner
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# TwIL-LM3 — Official Demo
+
+Official Hugging Face Space for
+[`webAI-Official/TwIL-LM3`](https://huggingface.co/webAI-Official/TwIL-LM3),
+webAI's 3B reasoning model for **formal logic**.
+
+TwIL-LM3 is built from
+[`HuggingFaceTB/SmolLM3-3B`](https://huggingface.co/HuggingFaceTB/SmolLM3-3B)
+through LoRA SFT, checkpoint fusion, WiSE-FT interpolation, and entropy-weighted
+GRPO. It emits a collapsible `<think>…</think>` trace before the answer.
+
+**Temperature = 0** (greedy) matches the published evaluation protocol. Keep
+max new tokens at 2048 or higher so the reasoning block is not truncated.
+
+This Space runs on ZeroGPU and streams tokens as they are generated.
+
+> **License:** the model is released under the *webAI Non-Commercial License
+> ver. 1.0* — see the [model repository](https://huggingface.co/webAI-Official/TwIL-LM3)
+> for full terms.

app.py

@@ -0,0 +1,235 @@
+import os
+import re
+from collections.abc import Iterator
+from threading import Thread
+
+import spaces
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+MODEL_ID = "webAI-Official/TwIL-LM3"
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))
+
+# Reasoning delimiters are non-special tokens (ids 128002 / 128003), so they
+# survive skip_special_tokens=True and Gradio can render them as a collapsible
+# section via allow_tags=["think"]. Built with chr() so the literal tags are
+# not written into the source.
+THINK_OPEN = chr(60) + "think" + chr(62)
+THINK_CLOSE = chr(60) + chr(47) + "think" + chr(62)
+_THINK_RE = re.compile(
+    re.escape(THINK_OPEN) + r".*?(" + re.escape(THINK_CLOSE) + r"|$)",
+    re.DOTALL,
+)
+
+TITLE = "# TwIL-LM3"
+DESCRIPTION = f"""
+Official demo of **[TwIL-LM3](https://huggingface.co/{MODEL_ID})**, webAI's 3B
+reasoning model for formal logic — FOL translation, entailment, semantic
+parsing, Lean formalisation and critique.
+
+Built from [SmolLM3-3B](https://huggingface.co/HuggingFaceTB/SmolLM3-3B) via
+LoRA SFT, checkpoint fusion, WiSE-FT, and entropy-weighted GRPO. It writes a
+collapsible `{THINK_OPEN}` reasoning trace before the answer.
+
+**Temperature = 0** (greedy) reproduces the published numbers. Keep max new
+tokens at **2048+** so the thinking block is not cut off.
+"""
+
+FOOTER = f"""
+---
+Official Space for [{MODEL_ID}](https://huggingface.co/{MODEL_ID}) ·
+License: [webAI Non-Commercial License ver. 1.0](https://huggingface.co/{MODEL_ID})
+"""
+
+PLACEHOLDER = """
+<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
+  <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">TwIL-LM3</h1>
+  <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Ask a formal-logic question…</p>
+</div>
+"""
+
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+).to("cuda")
+model.eval()
+model.generation_config.use_cache = True
+
+EOS_TOKEN_ID = tokenizer.eos_token_id
+
+
+def _strip_thinking(text: str) -> str:
+    """Drop the reasoning trace from a prior assistant turn before re-prompting."""
+    if not isinstance(text, str):
+        return ""
+    return _THINK_RE.sub("", text).strip()
+
+
+def _gpu_seconds(
+    message,
+    history,
+    max_new_tokens=2048,
+    temperature=0,
+    top_p=0.95,
+    enable_thinking=True,
+    *args,
+    **kwargs,
+):
+    tokens = int(max_new_tokens or 2048)
+    return min(180, max(45, 25 + tokens // 18))
+
+
+@spaces.GPU(duration=_gpu_seconds)
+def chat_twil_lm3(
+    message: str,
+    history: list,
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    enable_thinking: bool,
+) -> Iterator[str]:
+    conversation = []
+    for msg in history or []:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        if isinstance(content, list):
+            content = ""
+        if role == "assistant":
+            content = _strip_thinking(content)
+        conversation.append({"role": role, "content": content})
+    conversation.append({"role": "user", "content": message})
+
+    encoded = tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+        enable_thinking=enable_thinking,
+    )
+    input_ids = encoded["input_ids"]
+    attention_mask = encoded["attention_mask"]
+    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+        attention_mask = attention_mask[:, -MAX_INPUT_TOKEN_LENGTH:]
+        gr.Warning(
+            f"Trimmed the conversation to the last {MAX_INPUT_TOKEN_LENGTH} tokens."
+        )
+
+    input_ids = input_ids.to(model.device)
+    attention_mask = attention_mask.to(model.device)
+
+    streamer = TextIteratorStreamer(
+        tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        streamer=streamer,
+        max_new_tokens=int(max_new_tokens),
+        num_beams=1,
+        use_cache=True,
+        eos_token_id=EOS_TOKEN_ID,
+        pad_token_id=tokenizer.pad_token_id or EOS_TOKEN_ID,
+    )
+    if temperature == 0:
+        generate_kwargs["do_sample"] = False
+    else:
+        generate_kwargs["do_sample"] = True
+        generate_kwargs["temperature"] = float(temperature)
+        generate_kwargs["top_p"] = float(top_p)
+
+    Thread(target=model.generate, kwargs=generate_kwargs, daemon=True).start()
+
+    chunks: list[str] = []
+    for text in streamer:
+        chunks.append(text)
+        yield "".join(chunks)
+
+
+chatbot = gr.Chatbot(
+    height=450,
+    placeholder=PLACEHOLDER,
+    label="TwIL-LM3",
+    allow_tags=["think"],
+    line_breaks=False,
+)
+
+with gr.Blocks(theme=gr.themes.Soft(), css=CSS, fill_height=True) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(TITLE)
+        gr.Markdown(DESCRIPTION)
+        gr.ChatInterface(
+            fn=chat_twil_lm3,
+            chatbot=chatbot,
+            fill_height=True,
+            concurrency_limit=1,
+            additional_inputs_accordion=gr.Accordion(
+                label="Parameters", open=False, render=False
+            ),
+            additional_inputs=[
+                gr.Slider(
+                    minimum=256,
+                    maximum=4096,
+                    step=256,
+                    value=2048,
+                    label="Max new tokens",
+                    render=False,
+                    info="Keep this at 2048+ so the reasoning trace is not truncated.",
+                ),
+                gr.Slider(
+                    minimum=0,
+                    maximum=1.5,
+                    step=0.05,
+                    value=0,
+                    label="Temperature (0 = greedy)",
+                    render=False,
+                    info="Greedy (0) matches the published evaluation.",
+                ),
+                gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.05,
+                    value=0.95,
+                    label="Top-p",
+                    render=False,
+                    info="Used only when temperature > 0.",
+                ),
+                gr.Checkbox(
+                    value=True,
+                    label="Enable thinking (reasoning trace)",
+                    render=False,
+                    info="When on, the model reasons in a hidden block before answering.",
+                ),
+            ],
+            examples=[
+                [
+                    "Does 'All dogs are mammals. Rex is a dog.' entail 'Rex is a mammal'? "
+                    "Answer entailment, contradiction, or neutral."
+                ],
+                [
+                    "Translate into first-order logic: Every student who studies hard passes at least one exam."
+                ],
+                [
+                    "Formalize in Lean 4: If n is even, then n^2 is even."
+                ],
+                [
+                    "Is 'All birds fly. Tweety is a bird. Therefore Tweety flies.' logically valid? Explain."
+                ],
+                [
+                    "Formalize and evaluate: If it rains, the ground is wet. The ground is not wet. Therefore it did not rain."
+                ],
+            ],
+            cache_examples=False,
+        )
+        gr.Markdown(FOOTER)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+transformers>=5.0.0
+accelerate