""" Axiom-Ref — HuggingFace Space / Gradio App Governed Language Model: every output ships its own proof. """ 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 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-ref/mdlm_best.pt", weights_only=True, map_location=device)) tokenizer = Tokenizer.from_file("models/axiom-ref/bpe_tokenizer.json") bpe_vocab = tokenizer.get_vocab_size() BPE_BOS = tokenizer.token_to_id("") BPE_EOS = tokenizer.token_to_id("") 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) _dec_state = torch.load("models/axiom-ref/decoder_best.pt", weights_only=True, map_location=device) # Remap legacy weight names _dec_state = {k.replace("triad_embedding", "struct_embedding").replace("triad_pos", "struct_pos"): v for k, v in _dec_state.items()} decoder.load_state_dict(_dec_state) decoder.eval() def generate_governed(num_candidates=10, temperature=0.7): """Run the full 4-phase governed pipeline.""" # Phase 1: PROPOSE candidates = propose(mdlm, num_candidates=num_candidates, g_slots=2, s_slots=2, f_slots=2) # Phase 2: DECIDE 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] # Phase 3: PROMOTE promoted = promote(admitted) if not promoted: return "No candidates passed governance.", "", "{}", "" # Phase 4: EXECUTE outputs = execute(promoted) example, commitment = promoted[0] gov_dict = outputs[0].gov_structure # Generate prose 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) # Build governance trace output_hash = sha256(prose.encode()).hexdigest() gate_html = "" gate_names = ["G1 Structural Integrity", "G2 Completeness", "G3 Witness Sufficiency", "G4 Authority Separation", "G5 Provenance Continuity", "G6 Semantic Stability", "G7 Behavioral Prediction"] for g in gate_names: gate_html += f'
PASS {g}
' witness_html = "" for w_name, w_data in commitment.witnesses.items(): status = "ATTESTED" if w_data["attested"] else "WITHHELD" color = "#4ade80" if w_data["attested"] else "#e94560" witness_html += f'
{status} {w_name}
' trace = { "output_hash": output_hash[:32] + "...", "commitment": commitment.witness_bundle_hash[:32] + "...", "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_attested": 7, "admission": f"{t_count}/{num_candidates}", "timestamp": datetime.now(timezone.utc).isoformat(), } stats_html = f"""
PIPELINE STATS
Proposed: {num_candidates} | Admitted: {t_count} | Rejected: {f_count}
GATES
{gate_html}
WITNESSES
{witness_html}
COMMITMENT
{commitment.witness_bundle_hash[:48]}...
Irrevocable
""" return prose, stats_html, json.dumps(trace, indent=2), f"G: {trace['gov_structure']['G']}\nS: {trace['gov_structure']['S']}\nF: {trace['gov_structure']['F']}" # ── Gradio Interface ── with gr.Blocks( title="Axiom-Ref: Governed Language Model", theme=gr.themes.Base(primary_hue="green", neutral_hue="slate"), css=""" .output-prose { font-family: 'Courier New', monospace; font-size: 14px; } """ ) as app: gr.Markdown(""" # Axiom-Ref **Governed Language Model — every output ships its own proof.** Four phases: PROPOSE → DECIDE → PROMOTE → EXECUTE. No other language model ships a machine-verifiable governance trace with its output. """) 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="Decoder temperature") generate_btn = gr.Button("Generate Governed Output", variant="primary", size="lg") gr.Markdown("### Generated Output") output_prose = gr.Code(label="Governed Prose", language="c", lines=12) output_structure = gr.Textbox(label="Governed Structure", lines=3) with gr.Column(scale=1): gr.Markdown("### Governance Trace") governance_panel = gr.HTML() trace_json = gr.Code(label="Machine-Verifiable Trace (JSON)", language="json", lines=15) generate_btn.click( fn=generate_governed, inputs=[num_candidates, temperature], outputs=[output_prose, governance_panel, trace_json, output_structure], ) gr.Markdown(""" --- *[MetaCortex Dynamics DAO](https://github.com/MetaCortex-Dynamics) · [Source](https://github.com/MetaCortex-Dynamics/Axiom-Ref) · MIT License* """) if __name__ == "__main__": app.launch()