| import spaces |
|
|
| import asyncio |
| import json |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import gradio as gr |
| import torch |
| from huggingface_hub import snapshot_download |
| from safetensors.torch import load_file as load_safetensors |
| from transformers import AutoTokenizer |
|
|
| from cid.accelerator import wrap_torch_autocast |
| from cid.data import TrajectoryExample |
| from cid.defaults import ( |
| DEFAULT_ALLOCATION_THRESHOLD, |
| DEFAULT_BINDING_THRESHOLD, |
| DEFAULT_CONVERGENCE_THRESHOLD, |
| DEFAULT_NEED_THRESHOLD, |
| ) |
| from cid.model import ( |
| CIDMaterializerConfig, |
| ILLaDACIDConfig, |
| load_cid_adapter_from_pretrained, |
| load_cid_adapter_parameter_state, |
| ) |
| from cid.model.benchmark import run_neural_benchmark_case |
| from cid.model.encoding import ILLaDATextEncoder |
| from cid.runtime.engine import RuntimeConfig |
|
|
| MODEL_ID = "fwerkor/CID-v1-0.4B" |
| MODEL_REVISION = "b291b799cf654f7074e6bebb48d49fee31c17824" |
| MODEL_CHOICES = [("CID-v1-0.4B · 398.8M", MODEL_ID)] |
| DEMO_CODE_REVISION = "dffbec3" |
| MAX_PROMPT_CHARS = 4_000 |
| MAX_WORKSPACE_CHARS = 24_000 |
| MAX_DOCUMENTS = 12 |
|
|
| CUSTOM_CSS = r""" |
| .gradio-container { |
| max-width: 1280px !important; |
| margin: 0 auto !important; |
| padding: 28px 24px 36px !important; |
| } |
| |
| #cid-hero { |
| padding: 10px 2px 24px; |
| margin-bottom: 20px; |
| border-bottom: 1px solid var(--border-color-primary); |
| } |
| |
| #cid-hero .hero-kicker { |
| margin: 0 0 10px; |
| color: var(--primary-600); |
| font-size: 12px; |
| font-weight: 750; |
| letter-spacing: .11em; |
| text-transform: uppercase; |
| } |
| |
| #cid-hero .hero-title { |
| margin: 0; |
| max-width: 980px; |
| font-size: clamp(32px, 5vw, 54px); |
| font-weight: 760; |
| line-height: 1.04; |
| letter-spacing: -.035em; |
| } |
| |
| #cid-hero .hero-copy { |
| max-width: 900px; |
| margin: 14px 0 0; |
| color: var(--body-text-color-subdued); |
| font-size: 16px; |
| line-height: 1.65; |
| } |
| |
| #cid-hero .hero-links { |
| display: flex; |
| flex-wrap: wrap; |
| gap: 8px; |
| margin-top: 18px; |
| } |
| |
| #cid-hero .hero-links a { |
| display: inline-flex; |
| align-items: center; |
| min-height: 32px; |
| padding: 5px 10px; |
| border: 1px solid var(--border-color-primary); |
| border-radius: 999px; |
| color: var(--body-text-color); |
| background: var(--block-background-fill); |
| font-size: 13px; |
| font-weight: 600; |
| text-decoration: none; |
| } |
| |
| #cid-hero .hero-links a:hover { |
| border-color: var(--primary-400); |
| color: var(--primary-600); |
| } |
| |
| #prompt-panel, #result-panel { |
| padding: 18px; |
| border: 1px solid var(--border-color-primary); |
| border-radius: 18px; |
| background: var(--block-background-fill); |
| } |
| |
| #prompt-panel h3, #result-panel h3 { |
| margin-top: 0; |
| } |
| |
| #model-selector { |
| margin-bottom: 4px; |
| } |
| |
| #model-selector label { |
| font-weight: 650; |
| } |
| |
| #cid-prompt textarea { |
| min-height: 112px !important; |
| font-size: 15px !important; |
| line-height: 1.55 !important; |
| } |
| |
| #cid-workspace textarea { |
| overflow-y: auto !important; |
| scrollbar-gutter: stable; |
| overscroll-behavior: contain; |
| } |
| |
| #final-display textarea { |
| min-height: 176px !important; |
| font-size: 21px !important; |
| font-weight: 650 !important; |
| line-height: 1.5 !important; |
| } |
| |
| #run-button { |
| min-height: 48px; |
| margin-top: 4px; |
| font-weight: 700; |
| } |
| |
| #run-stats .stats { |
| display: grid; |
| grid-template-columns: repeat(4, minmax(0, 1fr)); |
| gap: 8px; |
| margin-top: 10px; |
| } |
| |
| #run-stats .stat { |
| padding: 10px 11px; |
| border: 1px solid var(--border-color-primary); |
| border-radius: 12px; |
| background: var(--background-fill-secondary); |
| } |
| |
| #run-stats .stat-value { |
| display: block; |
| color: var(--body-text-color); |
| font-size: 16px; |
| font-weight: 720; |
| line-height: 1.2; |
| } |
| |
| #run-stats .stat-label { |
| display: block; |
| margin-top: 3px; |
| color: var(--body-text-color-subdued); |
| font-size: 11px; |
| line-height: 1.25; |
| } |
| |
| #examples { |
| margin-top: 14px; |
| } |
| |
| #inspection { |
| margin-top: 22px; |
| } |
| |
| #inspection .tab-nav { |
| border-bottom-color: var(--border-color-primary); |
| } |
| |
| #runtime-trace table { |
| font-size: 12px !important; |
| } |
| |
| #cid-footer { |
| margin-top: 28px; |
| padding-top: 16px; |
| border-top: 1px solid var(--border-color-primary); |
| color: var(--body-text-color-subdued); |
| font-size: 13px; |
| } |
| |
| #cid-footer a { |
| color: var(--body-text-color); |
| font-weight: 600; |
| text-decoration: none; |
| } |
| |
| #cid-footer a:hover { |
| color: var(--primary-600); |
| } |
| |
| @media (max-width: 820px) { |
| .gradio-container { |
| padding: 18px 12px 28px !important; |
| } |
| |
| #cid-hero { |
| padding-top: 2px; |
| } |
| |
| #prompt-panel, #result-panel { |
| padding: 14px; |
| border-radius: 14px; |
| } |
| |
| #run-stats .stats { |
| grid-template-columns: repeat(2, minmax(0, 1fr)); |
| } |
| } |
| """ |
|
|
| WORKSPACE_DESCRIPTORS = ( |
| { |
| "name": "workspace_search", |
| "description": ( |
| "Search a task-local hidden document workspace and return candidate resource IDs " |
| "and titles. The workspace may be irrelevant to the task." |
| ), |
| "arguments": ({"name": "query", "kind": "string", "required": True},), |
| "cacheable": True, |
| "dynamic": False, |
| "versioned": False, |
| }, |
| { |
| "name": "workspace_read", |
| "description": "Read one task-local hidden resource by resource_id.", |
| "arguments": ({"name": "resource_id", "kind": "string", "required": True},), |
| "cacheable": True, |
| "dynamic": False, |
| "versioned": False, |
| }, |
| ) |
|
|
|
|
| def _load_model_bundle(): |
| model_dir = Path( |
| snapshot_download( |
| MODEL_ID, |
| revision=MODEL_REVISION, |
| allow_patterns=[ |
| "*.json", |
| "*.jinja", |
| "*.safetensors", |
| "*.pt", |
| "*.py", |
| ], |
| ) |
| ) |
| release_config = json.loads((model_dir / "cid_config.json").read_text(encoding="utf-8")) |
| adapter_config = ILLaDACIDConfig(**release_config["adapter_config"]) |
| tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) |
|
|
| device = torch.device("cuda") |
| adapter = load_cid_adapter_from_pretrained( |
| str(model_dir), |
| config=adapter_config, |
| freeze_backbone=True, |
| dtype=torch.float32, |
| low_cpu_mem_usage=True, |
| ).to(device) |
| adapter.set_backbone_trainable(True) |
| load_cid_adapter_parameter_state( |
| adapter, |
| load_safetensors(str(model_dir / "cid_adapter.safetensors"), device="cpu"), |
| ) |
|
|
| semantic_state = torch.load( |
| model_dir / "semantic-embedding.pt", |
| map_location="cpu", |
| weights_only=False, |
| ) |
| text_encoder = ILLaDATextEncoder.from_frozen_snapshot_state( |
| adapter, |
| tokenizer, |
| semantic_state, |
| device=device, |
| embedding_device="cpu", |
| ) |
| forward_model = wrap_torch_autocast( |
| torch, |
| adapter, |
| device_type="cuda", |
| dtype=torch.bfloat16, |
| ) |
| adapter.eval() |
| forward_model.eval() |
| return adapter, tokenizer, text_encoder, forward_model |
|
|
|
|
| ADAPTER, TOKENIZER, TEXT_ENCODER, FORWARD_MODEL = _load_model_bundle() |
|
|
|
|
| def _parse_workspace(raw: str) -> tuple[dict[str, Any], ...]: |
| raw = (raw or "").strip() |
| if not raw: |
| return () |
| if len(raw) > MAX_WORKSPACE_CHARS: |
| raise gr.Error(f"Workspace is limited to {MAX_WORKSPACE_CHARS:,} characters.") |
|
|
| blocks = [block.strip() for block in raw.replace("\r\n", "\n").split("\n---\n")] |
| documents: list[dict[str, Any]] = [] |
| for index, block in enumerate(block for block in blocks if block): |
| if len(documents) >= MAX_DOCUMENTS: |
| break |
| lines = [line.strip() for line in block.splitlines() if line.strip()] |
| if not lines: |
| continue |
| first = lines[0] |
| if first.lower().startswith("title:"): |
| title = first.split(":", 1)[1].strip() |
| body_lines = lines[1:] |
| elif first.startswith("#"): |
| title = first.lstrip("#").strip() |
| body_lines = lines[1:] |
| else: |
| title = f"Document {index + 1}" |
| body_lines = lines |
| body = " ".join(body_lines).strip() |
| if not body: |
| body = title |
| sentences = [piece.strip() for piece in body.replace("!", ".").replace("?", ".").split(".")] |
| sentences = [piece for piece in sentences if piece] |
| documents.append( |
| { |
| "resource_id": f"doc-{index:02d}", |
| "title": title or f"Document {index + 1}", |
| "sentences": sentences or [body], |
| } |
| ) |
| return tuple(documents) |
|
|
|
|
| def _build_example(prompt: str, workspace: str) -> tuple[TrajectoryExample, int]: |
| prompt = (prompt or "").strip() |
| if not prompt: |
| raise gr.Error("Enter a prompt first.") |
| if len(prompt) > MAX_PROMPT_CHARS: |
| raise gr.Error(f"Prompt is limited to {MAX_PROMPT_CHARS:,} characters.") |
|
|
| documents = _parse_workspace(workspace) |
| descriptors = WORKSPACE_DESCRIPTORS if documents else () |
| metadata: dict[str, Any] = {} |
| if documents: |
| metadata = { |
| "benchmark_workspace_documents": list(documents), |
| "benchmark_workspace_search_top_k": min(5, len(documents)), |
| "benchmark_workspace_latency_steps": 2, |
| "interaction_pattern": "retrieval_qa", |
| } |
| example = TrajectoryExample( |
| example_id="space-demo", |
| prompt=prompt, |
| target_display="CID interactive demo", |
| source_descriptors=descriptors, |
| metadata=metadata, |
| ) |
| return example, len(documents) |
|
|
|
|
| def _trace_rows(events: tuple[dict[str, Any], ...]) -> list[list[Any]]: |
| rows: list[list[Any]] = [] |
| for event in events: |
| payload = dict(event.get("payload", {})) |
| display_text = payload.pop("display_materialized_text", None) |
| payload.pop("display_text", None) |
| payload.pop("display_token_ids", None) |
| payload.pop("display_visible_token_ids", None) |
| detail = json.dumps(payload, ensure_ascii=False, default=str) |
| if display_text: |
| detail = f"display={display_text!r} " + detail |
| rows.append( |
| [ |
| int(event.get("step", 0)), |
| round(float(event.get("timestamp_s", 0.0)), 4), |
| str(event.get("kind", "")), |
| detail[:1200], |
| ] |
| ) |
| return rows |
|
|
|
|
| def _display_evolution(events: tuple[dict[str, Any], ...]) -> str: |
| changes: list[str] = [] |
| previous: str | None = None |
| seen_step = False |
| for event in events: |
| if event.get("kind") != "model_step_finished": |
| continue |
| payload = event.get("payload", {}) |
| current = str(payload.get("display_materialized_text", "")).strip() |
| if seen_step and current == previous: |
| continue |
| visible = current or "(empty Display)" |
| changes.append(f"**Step {event.get('step', 0)}** — {visible}") |
| previous = current |
| seen_step = True |
| if not changes: |
| return "No Display states were recorded before the run ended." |
| return "\n\n".join(changes) |
|
|
|
|
| def _summary_html(summary: dict[str, Any]) -> str: |
| status = "Converged" if summary["converged"] else "Not converged" |
| return f""" |
| <div class="stats"> |
| <div class="stat"><span class="stat-value">{status}</span><span class="stat-label">run status</span></div> |
| <div class="stat"><span class="stat-value">{summary['runtime_steps']}</span><span class="stat-label">runtime steps</span></div> |
| <div class="stat"><span class="stat-value">{summary['wall_time_s']:.2f}s</span><span class="stat-label">wall time</span></div> |
| <div class="stat"><span class="stat-value">{summary['information_need_events']}</span><span class="stat-label">information needs</span></div> |
| </div> |
| """ |
|
|
|
|
| @spaces.GPU(duration=20) |
| def run_cid( |
| prompt: str, |
| workspace: str, |
| max_steps: int, |
| need_threshold: float, |
| convergence_threshold: float, |
| allocation_threshold: float, |
| binding_threshold: float, |
| model_id: str, |
| ): |
| if model_id != MODEL_ID: |
| raise gr.Error(f"Model {model_id!r} is not available in this Space yet.") |
|
|
| example, document_count = _build_example(prompt, workspace) |
| materializer_config = CIDMaterializerConfig( |
| need_threshold=float(need_threshold), |
| convergence_threshold=float(convergence_threshold), |
| allocation_threshold=float(allocation_threshold), |
| ) |
| runtime_config = RuntimeConfig( |
| max_steps=256, |
| max_total_steps=int(max_steps), |
| max_wall_time_s=55.0, |
| binding_threshold=float(binding_threshold), |
| trace_display=True, |
| ) |
|
|
| started = time.perf_counter() |
| result = asyncio.run( |
| run_neural_benchmark_case( |
| ADAPTER, |
| TOKENIZER, |
| example, |
| text_encoder=TEXT_ENCODER, |
| forward_model=FORWARD_MODEL, |
| seed_teacher_state=False, |
| denoising_steps=8, |
| materializer_config=materializer_config, |
| runtime_config=runtime_config, |
| seed=0, |
| ) |
| ) |
| elapsed = time.perf_counter() - started |
| events = result.trace_events |
| need_events = [event for event in events if event.get("kind") == "information_need"] |
| tool_related = [ |
| event |
| for event in events |
| if event.get("kind") |
| in { |
| "information_need", |
| "binding_active", |
| "job_started", |
| "job_finished", |
| "observation_available", |
| "cache_hit", |
| "quiescence_started", |
| "quiescence_resumed", |
| } |
| ] |
| summary = { |
| "model": model_id, |
| "model_revision": MODEL_REVISION[:12], |
| "runtime_code_revision": DEMO_CODE_REVISION, |
| "runtime_steps": result.runtime_steps, |
| "converged": bool(result.evaluation.converged), |
| "wall_time_s": round(elapsed, 3), |
| "workspace_documents": document_count, |
| "information_need_events": len(need_events), |
| "tool_related_events": len(tool_related), |
| "thresholds": { |
| "need": float(need_threshold), |
| "convergence": float(convergence_threshold), |
| "allocation": float(allocation_threshold), |
| "binding": float(binding_threshold), |
| }, |
| } |
| return ( |
| result.final_text or "(empty Display)", |
| _summary_html(summary), |
| _display_evolution(events), |
| _trace_rows(events), |
| ) |
|
|
|
|
|
|
| CYMBIDIUM_WORKSPACE = """Title: Patrinia |
| Patrinia is a genus of herbaceous plants in the honeysuckle family. |
| There are about 17 species native to grassy mountain habitats in China, Siberia and Japan. |
| These are unassuming clump-forming perennial plants having thin, erect stems with few leaves and bearing a terminal inflorescence with yellow or white flowers. |
| The use for this plant is to provide a flower through long hot summers. |
| --- |
| Title: Cymbidium |
| Cymbidium , or boat orchid, is a genus of 52 evergreen species in the orchid family Orchidaceae. |
| The new Latin genus name is derived from the Latin "cymba" meaning boat. |
| Its first known use was in 1815.""" |
|
|
| ADORABLE_WORKSPACE = """Title: Adorable (band) |
| Adorable was an alternative rock band, formed in Coventry in 1990. |
| The band consisted of band members Pete Fijalkowski (vocals, guitar), Robert Dillam (guitar), Stephen 'Wil' Williams (bass) and Kevin Gritton (drums). |
| --- |
| Title: Lit (band) |
| Lit is an American rock band, formed in 1995 in Fullerton, California. |
| They are best known for their hit song "My Own Worst Enemy".""" |
|
|
| EXAMPLES = [ |
| ["How many days are in one week?", ""], |
| ["Which genus has more species, Cymbidium or Patrinia?", CYMBIDIUM_WORKSPACE], |
| ["Which band was formed first, Lit or Adorable?", ADORABLE_WORKSPACE], |
| ] |
|
|
| WORKSPACE_PLACEHOLDER = """Title: Mission note |
| Project Orion launches on October 17, 2032. |
| --- |
| Title: Logistics |
| The launch vehicle is Borealis.""" |
|
|
| with gr.Blocks( |
| title="CID · Interactive Demo", |
| css=CUSTOM_CSS, |
| theme=gr.themes.Soft( |
| primary_hue=gr.themes.colors.indigo, |
| neutral_hue=gr.themes.colors.slate, |
| ), |
| ) as demo: |
| gr.HTML( |
| """ |
| <div class="hero-kicker">Diffusion-native tool-augmented reasoning</div> |
| <h1 class="hero-title">Continuous Interaction Diffusion</h1> |
| <p class="hero-copy"> |
| Run the 398.8M-parameter CID reference checkpoint and inspect how its visible Display |
| changes across diffusion steps. CID maintains a latent Typed Cognitive Tensor (TCT), |
| can express information needs during generation, and revises output instead of committing |
| strictly left-to-right. |
| </p> |
| <div class="hero-links"> |
| <a href="https://huggingface.co/collections/fwerkor/cid" target="_blank">CID Collection</a> |
| <a href="https://huggingface.co/fwerkor/CID-v1-0.4B" target="_blank">CID-v1-0.4B model</a> |
| <a href="https://huggingface.co/papers/2608.10438" target="_blank">Paper on Hugging Face</a> |
| <a href="https://github.com/fwerkor/continuous-interaction-diffusion" target="_blank">Source code</a> |
| </div> |
| """, |
| elem_id="cid-hero", |
| ) |
|
|
| with gr.Row(equal_height=False): |
| with gr.Column(scale=5, min_width=360, elem_id="prompt-panel"): |
| gr.Markdown("### Try CID") |
| model_selector = gr.Dropdown( |
| choices=MODEL_CHOICES, |
| value=MODEL_ID, |
| label="Model", |
| interactive=True, |
| elem_id="model-selector", |
| ) |
| prompt = gr.Textbox( |
| label="Prompt", |
| value="What is 7 + 5?", |
| lines=4, |
| max_lines=8, |
| placeholder="Ask a short question...", |
| elem_id="cid-prompt", |
| ) |
|
|
| with gr.Accordion("Optional local workspace · experimental", open=False): |
| workspace = gr.Textbox( |
| label="Task-local documents", |
| value="", |
| lines=9, |
| max_lines=18, |
| placeholder=WORKSPACE_PLACEHOLDER, |
| info="Separate documents with a line containing ---. Start a block with 'Title:' or '#'. Retrieval is experimental on this compact checkpoint.", |
| elem_id="cid-workspace", |
| ) |
|
|
| gr.Examples( |
| examples=EXAMPLES, |
| inputs=[prompt, workspace], |
| label="Examples", |
| elem_id="examples", |
| ) |
|
|
| with gr.Accordion("Runtime controls", open=False): |
| max_steps = gr.Slider(1, 256, value=128, step=1, label="Max CID model steps") |
| with gr.Row(): |
| need_threshold = gr.Slider( |
| 0.05, |
| 0.95, |
| value=DEFAULT_NEED_THRESHOLD, |
| step=0.05, |
| label="Need threshold", |
| ) |
| convergence_threshold = gr.Slider( |
| 0.05, |
| 0.95, |
| value=DEFAULT_CONVERGENCE_THRESHOLD, |
| step=0.05, |
| label="Convergence threshold", |
| ) |
| with gr.Row(): |
| allocation_threshold = gr.Slider( |
| 0.05, |
| 0.95, |
| value=DEFAULT_ALLOCATION_THRESHOLD, |
| step=0.05, |
| label="Allocation threshold", |
| ) |
| binding_threshold = gr.Slider( |
| 0.05, |
| 0.95, |
| value=DEFAULT_BINDING_THRESHOLD, |
| step=0.05, |
| label="Binding threshold", |
| ) |
|
|
| run = gr.Button("Run CID", variant="primary", elem_id="run-button") |
|
|
| with gr.Column(scale=5, min_width=360, elem_id="result-panel"): |
| gr.Markdown("### Result") |
| final_display = gr.Textbox( |
| label="Final CID Display", |
| lines=7, |
| interactive=False, |
| elem_id="final-display", |
| ) |
| gr.Markdown("#### Run at a glance") |
| summary = gr.HTML( |
| '<div class="stats"><div class="stat"><span class="stat-value">Ready</span><span class="stat-label">run status</span></div><div class="stat"><span class="stat-value">—</span><span class="stat-label">runtime steps</span></div><div class="stat"><span class="stat-value">—</span><span class="stat-label">wall time</span></div><div class="stat"><span class="stat-value">—</span><span class="stat-label">information needs</span></div></div>', |
| elem_id="run-stats", |
| ) |
|
|
| with gr.Tabs(elem_id="inspection"): |
| with gr.Tab("Display evolution"): |
| gr.Markdown( |
| "The visible Display is materialized repeatedly; changes below expose revision across model steps." |
| ) |
| display_evolution = gr.Markdown() |
| with gr.Tab("Runtime trace"): |
| gr.Markdown( |
| "Low-level runtime events, including TCT lifecycle transitions and any tool-related activity." |
| ) |
| trace = gr.Dataframe( |
| headers=["step", "time_s", "event", "details"], |
| datatype=["number", "number", "str", "str"], |
| interactive=False, |
| wrap=True, |
| elem_id="runtime-trace", |
| ) |
|
|
| run.click( |
| fn=run_cid, |
| inputs=[ |
| prompt, |
| workspace, |
| max_steps, |
| need_threshold, |
| convergence_threshold, |
| allocation_threshold, |
| binding_threshold, |
| model_selector, |
| ], |
| outputs=[final_display, summary, display_evolution, trace], |
| concurrency_limit=1, |
| api_name="run_cid", |
| ) |
|
|
| gr.HTML( |
| """ |
| <div> |
| <strong>CID-v1-0.4B</strong> · 398.8M parameters · research prototype<br> |
| <a href="https://huggingface.co/collections/fwerkor/cid" target="_blank">Collection</a> |
| · |
| <a href="https://huggingface.co/fwerkor/CID-v1-0.4B" target="_blank">Model</a> |
| · |
| <a href="https://github.com/fwerkor/continuous-interaction-diffusion" target="_blank">Code</a> |
| · |
| <a href="https://huggingface.co/papers/2608.10438" target="_blank">Paper</a> |
| </div> |
| """, |
| elem_id="cid-footer", |
| ) |
|
|
| if __name__ == "__main__": |
| demo.queue(default_concurrency_limit=1).launch() |
|
|
