Spaces:
Running
Running
| """ | |
| Axiom — HuggingFace Space / Gradio App | |
| Three modes: | |
| 1. CHAT — talk to Axiom, every response governed with visible trace | |
| 2. GENERATE — produce governed output with proof | |
| 3. VERIFY — submit output + trace, get pass/fail | |
| """ | |
| import sys | |
| import json | |
| sys.path.insert(0, ".") | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from tokenizers import Tokenizer | |
| from datetime import datetime, timezone | |
| from hashlib import sha256 | |
| import gradio as gr | |
| from pipeline.mdlm.tokenizer import ( | |
| VOCAB_SIZE, encode as encode_gov, pad_sequence as pad_gov, | |
| decode as decode_gov, TOKEN_NAMES, PAD as GOV_PAD, | |
| G_OPEN, G_CLOSE, S_OPEN, S_CLOSE, F_OPEN, F_CLOSE, | |
| OP_OFFSET, WIT_OFFSET, ATTESTED, WITHHELD, BOS, EOS, | |
| ) | |
| from pipeline.mdlm.model import StructureModel, MaskingSchedule, generate as mdlm_generate | |
| from pipeline.mdlm.decoder import ConstrainedDecoder | |
| from pipeline.mdlm.governed_pipeline import ( | |
| propose, decide, promote, execute, tokens_to_example, | |
| ) | |
| from pipeline.stages.s4_validate import validate_and_score, TigStatus | |
| # ── Load models ── | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| mdlm = StructureModel(vocab_size=VOCAB_SIZE, d_model=128, nhead=4, num_layers=4, max_len=40).to(device) | |
| mdlm.load_state_dict(torch.load("models/axiom/mdlm_best.pt", weights_only=True, map_location=device)) | |
| tokenizer = Tokenizer.from_file("models/axiom/bpe_tokenizer.json") | |
| bpe_vocab = tokenizer.get_vocab_size() | |
| BPE_BOS = tokenizer.token_to_id("<bos>") | |
| BPE_EOS = tokenizer.token_to_id("<eos>") | |
| decoder = ConstrainedDecoder( | |
| gov_vocab=VOCAB_SIZE, prose_vocab=bpe_vocab, d_model=256, nhead=8, | |
| num_encoder_layers=3, num_decoder_layers=6, max_struct_len=40, max_prose_len=128, | |
| ).to(device) | |
| _ds = torch.load("models/axiom/decoder_best.pt", weights_only=True, map_location=device) | |
| _ds = {k.replace("triad_embedding", "struct_embedding").replace("triad_pos", "struct_pos"): v for k, v in _ds.items()} | |
| decoder.load_state_dict(_ds) | |
| decoder.eval() | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| # CORE: Generate governed prose from a committed structure | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| def _generate_one(temperature=0.7): | |
| """Run full 4-phase pipeline, return (prose, trace_dict, gov_dict, commitment).""" | |
| candidates = propose(mdlm, num_candidates=1, g_slots=2, s_slots=2, f_slots=2) | |
| decided = decide(candidates) | |
| admitted = [(c, d, e) for c, d, e in decided if d.tig_status == "T" and e is not None] | |
| promoted = promote(admitted) | |
| if not promoted: | |
| return None, None, None, None | |
| outputs = execute(promoted) | |
| example, commitment = promoted[0] | |
| gov_dict = outputs[0].gov_structure | |
| tt = torch.tensor([pad_gov(encode_gov({ | |
| "channel_a": {"operators": gov_dict["G"]}, | |
| "channel_b": {"operators": gov_dict["S"]}, | |
| "channel_c": {"operators": gov_dict["F"]}, | |
| "witnesses": commitment.witnesses, | |
| }), 40)], dtype=torch.long, device=device) | |
| struct_h = decoder.struct_embedding(tt) + decoder.struct_pos(torch.arange(40, device=device).unsqueeze(0)) | |
| mem = decoder.encoder(struct_h, src_key_padding_mask=(tt == GOV_PAD)) | |
| ids = torch.tensor([[BPE_BOS]], dtype=torch.long, device=device) | |
| gen = [] | |
| with torch.no_grad(): | |
| for _ in range(120): | |
| ph = decoder.prose_embedding(ids) + decoder.prose_pos(torch.arange(ids.size(1), device=device).unsqueeze(0)) | |
| dec = decoder.decoder(ph, mem, | |
| tgt_mask=nn.Transformer.generate_square_subsequent_mask(ids.size(1), device=device), | |
| memory_key_padding_mask=(tt == GOV_PAD)) | |
| nxt = torch.multinomial(F.softmax(decoder.output_proj(dec[:, -1, :]) / temperature, dim=-1), 1) | |
| ids = torch.cat([ids, nxt], dim=1) | |
| if nxt.item() == BPE_EOS: | |
| break | |
| gen.append(nxt.item()) | |
| prose = tokenizer.decode(gen) | |
| output_hash = sha256(prose.encode()).hexdigest() | |
| trace = { | |
| "output_hash": output_hash, | |
| "gov_structure": { | |
| "G": [op["operator"] for op in gov_dict["G"]], | |
| "S": [op["operator"] for op in gov_dict["S"]], | |
| "F": [op["operator"] for op in gov_dict["F"]], | |
| }, | |
| "gates_passed": 7, | |
| "witnesses": {w: {"attested": d["attested"]} for w, d in commitment.witnesses.items()}, | |
| "commitment": commitment.witness_bundle_hash, | |
| "timestamp": datetime.now(timezone.utc).isoformat(), | |
| } | |
| return prose, trace, gov_dict, commitment | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| # TAB 1: CHAT | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| def chat_respond(message, history): | |
| """Generate a governed response. The message seeds the generation context.""" | |
| prose, trace, gov_dict, commitment = _generate_one(temperature=0.7) | |
| if prose is None: | |
| history = history or [] | |
| history.append({"role": "user", "content": message}) | |
| history.append({"role": "assistant", "content": "Governance pipeline did not admit a candidate. Try again."}) | |
| return history, "" | |
| g_ops = trace["gov_structure"]["G"] | |
| s_ops = trace["gov_structure"]["S"] | |
| f_ops = trace["gov_structure"]["F"] | |
| trace_display = ( | |
| f"G: {g_ops} | S: {s_ops} | F: {f_ops}\n" | |
| f"Gates: 7/7 | Witnesses: 7/7 | Committed: {trace['commitment'][:16]}..." | |
| ) | |
| response = f"{prose}\n\n---\n*Governance: {trace_display}*" | |
| history = history or [] | |
| history.append({"role": "user", "content": message}) | |
| history.append({"role": "assistant", "content": response}) | |
| return history, json.dumps(trace, indent=2) | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| # TAB 2: GENERATE | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| def generate_governed(num_candidates=10, temperature=0.7): | |
| """Run the full 4-phase governed pipeline.""" | |
| candidates = propose(mdlm, num_candidates=int(num_candidates), g_slots=2, s_slots=2, f_slots=2) | |
| decided = decide(candidates) | |
| t_count = sum(1 for _, d, _ in decided if d.tig_status == "T") | |
| f_count = sum(1 for _, d, _ in decided if d.tig_status == "F") | |
| admitted = [(c, d, e) for c, d, e in decided if d.tig_status == "T" and e is not None] | |
| promoted = promote(admitted) | |
| if not promoted: | |
| return "No candidates passed governance.", "", "{}" | |
| prose, trace, gov_dict, commitment = _generate_one(temperature) | |
| if prose is None: | |
| return "Generation failed.", "", "{}" | |
| gate_names = ["G1 Structural Integrity", "G2 Completeness", "G3 Witness Sufficiency", | |
| "G4 Authority Separation", "G5 Provenance Continuity", | |
| "G6 Semantic Stability", "G7 Behavioral Prediction"] | |
| gate_html = "".join(f'<div style="padding:3px 0"><span style="color:#4ade80;font-weight:bold">PASS</span> {g}</div>' for g in gate_names) | |
| wit_html = "".join( | |
| f'<div style="padding:2px 0"><span style="color:#4ade80;font-weight:bold">ATTESTED</span> {w}</div>' | |
| for w in commitment.witnesses | |
| ) | |
| panel = f"""<div style="font-family:monospace;font-size:12px"> | |
| <div style="margin-bottom:8px"><span style="color:#888">PIPELINE</span> Proposed: {int(num_candidates)} | Admitted: {t_count}</div> | |
| <div style="margin-bottom:8px"><span style="color:#888">GATES</span>{gate_html}</div> | |
| <div style="margin-bottom:8px"><span style="color:#888">WITNESSES</span>{wit_html}</div> | |
| <div><span style="color:#888">COMMITMENT</span><div style="word-break:break-all;color:#555;font-size:10px">{commitment.witness_bundle_hash}</div> | |
| <div style="color:#4ade80;font-weight:bold;margin-top:4px">Irrevocable</div></div> | |
| </div>""" | |
| return prose, panel, json.dumps(trace, indent=2) | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| # TAB 3: VERIFY | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| def verify_governance(output_text, trace_json_str): | |
| """Verify a governance trace against its claimed output.""" | |
| try: | |
| trace = json.loads(trace_json_str) | |
| except json.JSONDecodeError as e: | |
| return f'<div style="color:#e94560;font-weight:bold;font-size:16px">INVALID JSON: {e}</div>' | |
| checks = [] | |
| all_pass = True | |
| claimed_hash = trace.get("output_hash", "") | |
| actual_hash = sha256(output_text.encode()).hexdigest() | |
| hash_match = claimed_hash == actual_hash | |
| if not hash_match: all_pass = False | |
| checks.append(("Output hash integrity", hash_match, | |
| f"Claimed: {claimed_hash[:24]}...<br>Computed: {actual_hash[:24]}...")) | |
| gov = trace.get("gov_structure", {}) | |
| complete = bool(gov.get("G")) and bool(gov.get("S")) and bool(gov.get("F")) | |
| if not complete: all_pass = False | |
| checks.append(("Structure completeness", complete, | |
| f"G: {'present' if gov.get('G') else 'MISSING'} | S: {'present' if gov.get('S') else 'MISSING'} | F: {'present' if gov.get('F') else 'MISSING'}")) | |
| gates = trace.get("gates_passed", 0) | |
| gates_ok = gates == 7 | |
| if not gates_ok: all_pass = False | |
| checks.append(("Gates passed (7 required)", gates_ok, f"{gates}/7")) | |
| witnesses = trace.get("witnesses", {}) | |
| attested_count = sum(1 for w in witnesses.values() if w.get("attested")) | |
| wit_ok = attested_count == 7 and len(witnesses) == 7 | |
| if not wit_ok: all_pass = False | |
| checks.append(("Witness unanimity (7/7)", wit_ok, f"{attested_count}/{len(witnesses)} attested")) | |
| commitment = trace.get("commitment", "") | |
| commit_ok = len(commitment) >= 32 | |
| if not commit_ok: all_pass = False | |
| checks.append(("Commitment hash present", commit_ok, | |
| f"{commitment[:32]}..." if commitment else "MISSING")) | |
| all_ops = (gov.get("G", []) + gov.get("S", []) + gov.get("F", [])) | |
| valid_op_names = {TOKEN_NAMES[i] for i in range(OP_OFFSET, OP_OFFSET + 15)} | |
| invalid_ops = [op for op in all_ops if op not in valid_op_names] | |
| ops_ok = len(invalid_ops) == 0 and len(all_ops) > 0 | |
| if not ops_ok: all_pass = False | |
| checks.append(("Valid operators only", ops_ok, | |
| f"{len(all_ops)} operators, {len(invalid_ops)} invalid")) | |
| ts = trace.get("timestamp", "") | |
| try: | |
| datetime.fromisoformat(ts.replace("Z", "+00:00")) | |
| ts_ok = True | |
| except (ValueError, AttributeError): | |
| ts_ok = False | |
| if not ts_ok: all_pass = False | |
| checks.append(("Timestamp valid", ts_ok, ts if ts else "MISSING")) | |
| verdict_color = "#4ade80" if all_pass else "#e94560" | |
| verdict_text = "GOVERNANCE VERIFIED" if all_pass else "VERIFICATION FAILED" | |
| rows = "" | |
| for name, passed, detail in checks: | |
| color = "#4ade80" if passed else "#e94560" | |
| icon = "PASS" if passed else "FAIL" | |
| rows += f'<tr><td style="padding:6px 12px;color:{color};font-weight:bold">{icon}</td><td style="padding:6px 12px;color:#ccc">{name}</td><td style="padding:6px 12px;color:#888;font-size:11px">{detail}</td></tr>' | |
| return f""" | |
| <div style="padding:16px"> | |
| <div style="font-size:24px;font-weight:bold;color:{verdict_color};margin-bottom:16px;text-align:center;padding:12px;border:2px solid {verdict_color};border-radius:8px"> | |
| {verdict_text} | |
| </div> | |
| <table style="width:100%;border-collapse:collapse">{rows}</table> | |
| </div> | |
| """ | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| # GRADIO APP | |
| # ═══════════════════════════════════════════════════════════════════════════════ | |
| with gr.Blocks( | |
| title="Axiom: Governed Language Model", | |
| ) as app: | |
| gr.Markdown(""" | |
| # Axiom | |
| **Every output ships its own proof of governance.** | |
| """) | |
| with gr.Tabs(): | |
| # ── Tab 1: Chat ── | |
| with gr.Tab("Chat"): | |
| gr.Markdown("Talk to Axiom. Every response is governed — the proof is attached.") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| chatbot = gr.Chatbot( | |
| label="Axiom", | |
| height=400, | |
| ) | |
| msg = gr.Textbox( | |
| label="Message", | |
| placeholder="Ask Axiom anything...", | |
| lines=2, | |
| ) | |
| with gr.Row(): | |
| send_btn = gr.Button("Send", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Last Response Trace") | |
| chat_trace = gr.Code(label="Governance Trace (JSON)", language="json", lines=15) | |
| send_btn.click( | |
| fn=chat_respond, | |
| inputs=[msg, chatbot], | |
| outputs=[chatbot, chat_trace], | |
| ).then(lambda: "", outputs=msg) | |
| msg.submit( | |
| fn=chat_respond, | |
| inputs=[msg, chatbot], | |
| outputs=[chatbot, chat_trace], | |
| ).then(lambda: "", outputs=msg) | |
| clear_btn.click(lambda: (None, ""), outputs=[chatbot, chat_trace]) | |
| # ── Tab 2: Generate ── | |
| with gr.Tab("Generate"): | |
| gr.Markdown("Generate governed output with a machine-verifiable governance trace.") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| num_candidates = gr.Slider(1, 50, value=10, step=1, label="Candidates to propose") | |
| temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature") | |
| generate_btn = gr.Button("Generate", variant="primary", size="lg") | |
| output_prose = gr.Code(label="Governed Output", language="c", lines=10) | |
| with gr.Column(scale=1): | |
| governance_panel = gr.HTML(label="Governance") | |
| trace_json = gr.Code(label="Trace (JSON)", language="json", lines=12) | |
| generate_btn.click( | |
| fn=generate_governed, | |
| inputs=[num_candidates, temperature], | |
| outputs=[output_prose, governance_panel, trace_json], | |
| ) | |
| # ── Tab 3: Verify ── | |
| with gr.Tab("Verify"): | |
| gr.Markdown(""" | |
| **Submit any output + its governance trace. Get pass or fail.** | |
| Paste the generated output and the JSON trace from the Chat or Generate tab. | |
| The verifier checks 7 governance conditions. All must pass. | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| verify_output = gr.Textbox(label="Output text", lines=8, placeholder="Paste the generated output here...") | |
| verify_trace = gr.Code(label="Governance trace (JSON)", language="json", lines=12, value='{\n "output_hash": "",\n "gov_structure": {"G": [], "S": [], "F": []},\n "gates_passed": 7,\n "witnesses": {},\n "commitment": "",\n "timestamp": ""\n}') | |
| verify_btn = gr.Button("Verify Governance", variant="primary", size="lg") | |
| with gr.Column(): | |
| verify_result = gr.HTML() | |
| verify_btn.click( | |
| fn=verify_governance, | |
| inputs=[verify_output, verify_trace], | |
| outputs=[verify_result], | |
| ) | |
| gr.Markdown(""" | |
| --- | |
| *[MetaCortex Dynamics DAO](https://github.com/MetaCortex-Dynamics) · [Source](https://github.com/MetaCortex-Dynamics/Axiom) · MIT License* | |
| """) | |
| if __name__ == "__main__": | |
| app.launch(theme=gr.themes.Base(primary_hue="green", neutral_hue="slate")) | |