Passage QA demo for amx-reasoning-v1-instruct

README.md

@@ -1,13 +1,43 @@
 ---
-title: Amx Reasoning V1 Qa Demo
-emoji: 📉
-colorFrom: yellow
-colorTo: gray
+title: amx-reasoning-v1-instruct QA
+emoji: 🧮
+colorFrom: green
+colorTo: red
 sdk: gradio
 sdk_version: 6.26.0
-python_version: '3.13'
 app_file: app.py
-pinned: false
+short_description: Passage QA with a 7.5M-param CPU-trained LM
+python_version: "3.12"
+startup_duration_timeout: 30m
+license: apache-2.0
+models:
+  - gdiamos/amx-reasoning-v1-instruct
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# amx-reasoning-v1-instruct — passage QA
+
+Demo for [`gdiamos/amx-reasoning-v1-instruct`](https://huggingface.co/gdiamos/amx-reasoning-v1-instruct):
+a 7,492,448-parameter causal LM (3,315,552 active per token) trained end to end
+on a single Intel Emerald Rapids CPU core. Give it a passage and a question and
+it extracts a one-or-two-word answer.
+
+The demo runs the model's own reference path, unmodified:
+
+- the `m2r` package that trained it (the architecture is not a `transformers`
+  one — `AutoModelForCausalLM` will not load it),
+- `render_prompt(..., thinking=False)` for the exact prompt format,
+- greedy argmax decoding, stopped on the `<SPECIAL_12>` end-of-turn token,
+- the 800-token vocabulary mask from `generation.json` applied before the
+  argmax (required for correct output; the demo exposes a toggle so you can see
+  what happens without it).
+
+Inference runs on CPU, which is the hardware the model was designed and trained
+for — a forward pass at these dimensions is a few GFLOP.
+
+## Example attribution
+
+Example passages are drawn from the datasets the model card reports on:
+[DROP](https://huggingface.co/datasets/ucinlp/drop) and
+[SQuAD v2](https://huggingface.co/datasets/rajpurkar/squad_v2) (CC BY-SA 4.0),
+and [databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k)
+(CC BY-SA 3.0). The first example is the model repo's own `example.py`.

app.py

@@ -0,0 +1,253 @@
+"""Gradio demo for gdiamos/amx-reasoning-v1-instruct.
+
+A 7.5M-parameter (3.3M active) causal LM trained end-to-end on a single Intel
+AMX CPU core. It does passage-grounded extractive QA: hand it a passage and a
+question, get a one-or-two-word answer back.
+
+The inference path here mirrors the model repo's own `example.py` exactly:
+the `m2r` package that trained it, `render_prompt(..., thinking=False)` for the
+prompt format, greedy argmax decoding stopped on EOT, and the 800-token
+vocabulary mask from `generation.json` applied before the argmax.
+"""
+
+import json
+import pathlib
+import sys
+import time
+
+import gradio as gr
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+MODEL_ID = "gdiamos/amx-reasoning-v1-instruct"
+
+# The model's own source ships in the repo under m2r/ -- it is not a
+# transformers architecture, so AutoModelForCausalLM will not load it.
+LOCAL = pathlib.Path(
+    snapshot_download(
+        MODEL_ID,
+        allow_patterns=[
+            "m2r/**",
+            "config.json",
+            "training_config.yaml",
+            "generation.json",
+            "tokenizer.json",
+            "model.safetensors",
+        ],
+    )
+)
+sys.path.insert(0, str(LOCAL))
+
+from m2r.config import load  # noqa: E402
+from m2r.data.templates import EOT, render_prompt  # noqa: E402
+from m2r.model.torch_model import Model, swa_mask  # noqa: E402
+
+torch.set_grad_enabled(False)
+
+cfg = load(LOCAL / "training_config.yaml")
+model = Model(cfg.model).to(torch.bfloat16)
+model.load_state_dict(load_file(str(LOCAL / "model.safetensors")))
+model.eval()
+
+tok = Tokenizer.from_file(str(LOCAL / "tokenizer.json"))
+MASK = swa_mask(cfg.model, dtype=torch.bfloat16)
+GEN = json.loads((LOCAL / "generation.json").read_text())
+
+# Required, not a knob: these 800 vocabulary rows never occur in the training
+# corpus, so the sampled softmax never drew them as negatives and never pushed
+# their logits down. They sit near 0 while trained-but-wrong tokens sit near
+# -7.9, so they win the argmax whenever the model is unsure. Before masking,
+# " ballo" and "Frequently" were 29% of all DROP answers.
+BAN = torch.tensor(GEN["banned_token_ids"], dtype=torch.long)
+
+PAD_TO = max(cfg.model.window, cfg.model.route_block or 1, 256)
+EOT_ID = tok.encode(EOT, add_special_tokens=False).ids[0]
+MAX_PASSAGE_TOKENS = 2000
+DEFAULT_MAX_NEW_TOKENS = int(GEN.get("recommended", {}).get("max_new_tokens", 32))
+
+EXAMPLES = json.loads((pathlib.Path(__file__).parent / "examples.json").read_text())
+
+print(
+    f"loaded {MODEL_ID}: "
+    f"{sum(p.numel() for p in model.parameters()):,} stored parameters, "
+    f"pad_to={PAD_TO}, eot_id={EOT_ID}, {len(BAN)} banned token ids",
+    flush=True,
+)
+
+
+def _clip_passage(passage: str) -> str:
+    """Trim a passage to MAX_PASSAGE_TOKENS on a token boundary."""
+    enc = tok.encode(passage, add_special_tokens=False)
+    if len(enc.ids) <= MAX_PASSAGE_TOKENS:
+        return passage
+    end = enc.offsets[MAX_PASSAGE_TOKENS - 1][1]
+    return passage[:end].rstrip() + " ..."
+
+
+def _next_token(ids: list[int], apply_vocab_mask: bool) -> int:
+    n = len(ids)
+    x = torch.tensor([ids + [0] * ((-n) % PAD_TO)])
+    h = model.body(x, MASK)[:, n - 1]
+    logits = (h @ model.emb.t().to(h.dtype)).float()[0]
+    if apply_vocab_mask:
+        logits[BAN] = -1e30
+    return int(logits.argmax())
+
+
+def answer(
+    question: str,
+    passage: str,
+    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
+    apply_vocab_mask: bool = True,
+) -> tuple[str, str]:
+    """Answer a question about a passage with amx-reasoning-v1-instruct.
+
+    Greedy decoding, stopped on the end-of-turn token, using the model's own
+    prompt template. The answer is normally one or two words extracted from
+    the passage.
+
+    Args:
+        question: the question to ask about the passage.
+        passage: the passage the answer should be grounded in.
+        max_new_tokens: hard cap on generated tokens (the model usually stops
+            after one or two).
+        apply_vocab_mask: apply the 800-token vocabulary mask from
+            generation.json. Required for correct output; turn it off only to
+            see the untrained-row failure mode the model card describes.
+
+    Returns:
+        The model's answer, and a one-line note about how it was produced.
+    """
+    question = (question or "").strip()
+    passage = (passage or "").strip()
+    if not question:
+        return "", "Enter a question."
+    if not passage:
+        return "", "This model is extractive — it needs a passage to answer from."
+
+    clipped = _clip_passage(passage)
+    prompt = render_prompt([f"{question}\n\n{clipped}"], thinking=False)
+    ids = tok.encode(prompt, add_special_tokens=False).ids
+    n_prompt = len(ids)
+
+    t0 = time.perf_counter()
+    out: list[int] = []
+    stopped = False
+    for _ in range(int(max_new_tokens)):
+        t = _next_token(ids, apply_vocab_mask)
+        if t == EOT_ID:
+            stopped = True
+            break
+        out.append(t)
+        ids.append(t)
+    dt = time.perf_counter() - t0
+
+    text = tok.decode(out).replace("<think></think>", "").strip()
+    if not text:
+        text = "(empty)"
+
+    note = (
+        f"{n_prompt} prompt tokens → {len(out)} generated in {dt:.2f}s on CPU "
+        f"({'stopped on EOT' if stopped else 'hit the token cap'})"
+    )
+    if not apply_vocab_mask:
+        note += " · **vocabulary mask off**"
+    if clipped is not passage:
+        note += f" · passage clipped to {MAX_PASSAGE_TOKENS} tokens"
+    return text, note
+
+
+CSS = """
+#col-container { max-width: 1040px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+#answer textarea { font-size: 20px; font-weight: 600; }
+"""
+
+INTRO = """# amx-reasoning-v1-instruct — passage QA
+
+A **7,492,448-parameter** causal LM (3,315,552 active per token) trained end to
+end on **one Intel Emerald Rapids CPU core**. Give it a passage and a question;
+it extracts a one-or-two-word answer and stops. It is a research artifact —
+18.2% exact match on held-out extractive QA — and the point is that a model
+this small does passage-grounded retrieval at all.
+
+It **retrieves and compares; it cannot calculate.** It also answers some
+unanswerable questions anyway. The examples below include those failures on
+purpose.
+
+[Model card](https://huggingface.co/gdiamos/amx-reasoning-v1-instruct) ·
+[Paper](https://huggingface.co/gdiamos/amx-reasoning-v1-instruct/blob/main/paper.pdf)
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="amx-reasoning-v1 QA") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(INTRO)
+
+        with gr.Row():
+            question = gr.Textbox(
+                label="Question",
+                placeholder="What year was the company founded?",
+                lines=1,
+                scale=4,
+            )
+            run = gr.Button("Answer", variant="primary", scale=1)
+
+        passage = gr.Textbox(
+            label="Passage",
+            placeholder="Paste the passage the answer should come from…",
+            lines=9,
+        )
+
+        answer_box = gr.Textbox(
+            label="Answer", elem_id="answer", lines=2, show_copy_button=True
+        )
+        info = gr.Markdown()
+
+        with gr.Accordion("Advanced", open=False):
+            max_new_tokens = gr.Slider(
+                1, 64, value=DEFAULT_MAX_NEW_TOKENS, step=1,
+                label="Max new tokens",
+                info="The model normally emits one or two and stops on EOT.",
+            )
+            apply_vocab_mask = gr.Checkbox(
+                value=True,
+                label="Apply the vocabulary mask from generation.json",
+                info=(
+                    "800 vocabulary rows never occurred in training, so their "
+                    "logits were never pushed down and they win the argmax "
+                    "whenever the model is unsure. Unchecking this reproduces "
+                    "the ' ballo' failure mode from the model card."
+                ),
+            )
+
+        gr.Examples(
+            examples=[[e["question"], e["passage"]] for e in EXAMPLES],
+            example_labels=[
+                f"{e['source']} — {e['question']}" for e in EXAMPLES
+            ],
+            inputs=[question, passage],
+            outputs=[answer_box, info],
+            fn=answer,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+        gr.Markdown(
+            "Example passages come from the datasets the model card reports on: "
+            "[DROP](https://huggingface.co/datasets/ucinlp/drop) and "
+            "[SQuAD v2](https://huggingface.co/datasets/rajpurkar/squad_v2) "
+            "(CC BY-SA 4.0), and "
+            "[databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) "
+            "(CC BY-SA 3.0). The first example is the model repo's own "
+            "`example.py`."
+        )
+
+    inputs = [question, passage, max_new_tokens, apply_vocab_mask]
+    outputs = [answer_box, info]
+    run.click(answer, inputs=inputs, outputs=outputs, api_name="answer")
+    question.submit(answer, inputs=inputs, outputs=outputs, api_name=False)
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

examples.json

@@ -0,0 +1,44 @@
+[
+ {
+  "question": "In what year was the Royal Dutch Petroleum Company founded?",
+  "passage": "In February 1907, the Royal Dutch Shell Group was created through the amalgamation of two rival companies: the Royal Dutch Petroleum Company, founded in 1890, and the Shell Transport and Trading Company of the United Kingdom.",
+  "source": "the model's own example.py",
+  "reference": "1890"
+ },
+ {
+  "question": "Which field goals did Neil Rackers make?",
+  "passage": "The Texans' fifteenth game was an AFC duel with the Broncos. The Texans commanded the first half with RB Arian Foster getting a 3-yard TD run, followed by QB Matt Schaub getting a 3-yard TD pass to TE Owen Daniels, then with kicker Neil Rackers hitting a 34-yard field goal. The Broncos got on the board with RB Correll Buckhalter getting a 3-yard TD run, but the Texans scored again with Rackers nailing a 54-yard field goal. The Broncos replied as kicker Steven Hauschka got a 27-yard field goal, but the Texans extended their lead with Rackers hitting a 57-yard field goal. However, they failed to maintain this lead after QB Tim Tebow completed a 23-yard TD pass to Buckhalter, followed by Tebow scrambling 6-yards for a touchdown.",
+  "source": "DROP",
+  "reference": "34-yard, 54-yard, 57-yard"
+ },
+ {
+  "question": "How many years did Manipur raid the Upper Chindwin region?",
+  "passage": "Manipur was a tributary to Burma in the 16th century  but had gone its own way since. It raided Upper Chindwin region in 1647 and 1692. However, in 1704, the raja of Manipur presented his daughter to Ava. Starting in the 1720s, Manipur under the leadership of Pamheiba  became a thorn to Upper Burma. In early 1724, the Manipuris raided Upper Burma. In response, an expedition force of 3,000 men marched to Manipur in November 1724. The army was ambushed in the swamps at Heirok, and retreated in haste. The Manipuris then returned ten years later. From 1735 to 1741, Manipuris raided the Upper Chindwin regions, increasingly deeper with each raid. Burmese defences were simply bypassed the Manipuris on their horseback. In December 1739, they reached as far as Sagaing, and looted and burned everything insight. The Burmese defences finally stopped them at Myedu in early 1741, with each side agreeing to an uneasy truce. But the Manipuris had annexed the Kabaw valley. The truce did not last. Another raid came all the way down to Ava in 1744. The last raid came in 1749. Upon arrival at Ava, the Manipuri chief found a large Burmese army, and presented his 12-year-old daughter instead, and left.",
+  "source": "DROP",
+  "reference": "46"
+ },
+ {
+  "question": "Where did researchers study chimps in heavily forested regions?",
+  "passage": "In 2006–07, researchers from the Wildlife Conservation Society studied gorillas in heavily forested regions centered on the Ouesso district of the Sangha Region. They suggest a population on the order of 125,000 Western Lowland Gorillas, whose isolation from humans has been largely preserved by inhospitable swamps.",
+  "source": "SQuAD v2 (unanswerable)",
+  "reference": "The passage does not say."
+ },
+ {
+  "question": "Which species avoided the polar areas?",
+  "passage": "The Early Cretaceous spans from 145 million to 100 million years ago. The Early Cretaceous saw the expansion of seaways, and as a result, the decline and extinction of sauropods (except in South America). Many coastal shallows were created, and that caused Ichthyosaurs to die out. Mosasaurs evolved to replace them as head of the seas. Some island-hopping dinosaurs, like Eustreptospondylus, evolved to cope with the coastal shallows and small islands of ancient Europe. Other dinosaurs rose up to fill the empty space that the Jurassic-Cretaceous extinction left behind, such as Carcharodontosaurus and Spinosaurus. Of the most successful would be the Iguanodon which spread to every continent. Seasons came back into effect and the poles got seasonally colder, but dinosaurs still inhabited this area like the Leaellynasaura which inhabited the polar forests year-round, and many dinosaurs migrated there during summer like Muttaburrasaurus. Since it was too cold for crocodiles, it was the last stronghold for large amphibians, like Koolasuchus. Pterosaurs got larger as species like Tapejara and Ornithocheirus evolved.",
+  "source": "SQuAD v2 (unanswerable)",
+  "reference": "The passage does not say."
+ },
+ {
+  "question": "What is the Kentucky Derby Trophy",
+  "passage": "The Kentucky Derby Trophy is a set of four trophies that are awarded to the winning connections of America's most famous race: the grade one $3,000,000 Kentucky Derby. The owner receives a gold trophy while the trainer, the jockey and the breeder win a silver half size replica of the main gold trophy. The trophy itself has been run for since the 50th running of the Kentucky Derby in 1924. Churchill Downs Race Course of Louisville, Kentucky has annually presented a gold trophy to the winning owner of the famed \"Run for the Roses.\"",
+  "source": "Dolly 15k",
+  "reference": "The Kentucky Derby Trophy is a set of four trophies that are awarded to the winning connections of America's most famous race: the grade one $3,000,000 Kentucky Derby. The owner receives a gold trophy while the trainer, the jockey and the breeder win a silver half size replica of the main gold trophy. The trophy itself has been run for since the 50th running of the Kentucky Derby in 1924. Churchill Downs Race Course of Louisville, Kentucky has annually presented a gold trophy to the winning owner of the famed \"Run for the Roses.\""
+ },
+ {
+  "question": "From the passage note down the name of the countries which have most voting power. List the results in comma separated format.",
+  "passage": "The World Bank is an international financial institution that provides loans and grants to the governments of low- and middle-income countries for the purpose of pursuing capital projects. The World Bank is the collective name for the International Bank for Reconstruction and Development (IBRD) and International Development Association (IDA), two of five international organizations owned by the World Bank Group. It was established along with the International Monetary Fund at the 1944 Bretton Woods Conference. After a slow start, its first loan was to France in 1947. In the 1970s, it focused on loans to developing world countries, shifting away from that mission in the 1980s. For the last 30 years, it has included NGOs and environmental groups in its loan portfolio. Its loan strategy is influenced by the United Nations' Sustainable Development Goals, as well as environmental and social safeguards.\n\nAs of 2022, the World Bank is run by a president and 25 executive directors, as well as 29 various vice presidents. IBRD and IDA have 189 and 174 member countries, respectively. The U.S., Japan, China, Germany and the U.K. have the most voting power. The bank aims loans at developing countries to help reduce poverty. The bank is engaged in several global partnerships and initiatives, and takes a role in working toward addressing climate change. The World Bank operates a number of training wings and it works with the Clean Air Initiative and the UN Development Business. It works within the Open Data Initiative and hosts an Open Knowledge Repository.\n\nThe World Bank has been criticized as promoting inflation and harming economic development, causing protests in 1988 and 2000. There has also been criticism of the bank's governance and response to the COVID-19 pandemic.",
+  "source": "Dolly 15k",
+  "reference": "U.S., Japan, China, Germany, U.K."
+ }
+]

requirements.txt

@@ -0,0 +1,4 @@
+torch
+safetensors
+tokenizers
+pyyaml