Support pipeline result JSON format

app.py

@@ -1,7 +1,9 @@
 """
 Argumentation Graph Visualizer — Hugging Face Space
-Supports .xml (Demosthenes) and .json (structured sections with inline tags).
-Graph is rendered in an iframe to allow full JavaScript execution.
+Supports:
+  - Demosthenes XML  (.xml)
+  - Structured JSON with inline <prem>/<conc> tags  (.json, format A)
+  - Pipeline result JSON with segment dicts          (.json, format B)
 """
 
 import re
@@ -24,7 +26,7 @@ def parse_xml(filepath):
         nid = elem.get("ID")
         if not nid:
             continue
-        nodes[nid] = elem.tag
+        nodes[nid] = elem.tag.lower()
         for t in filter(None, elem.get("SUP", "").split("|")):
             sup_edges.append((nid, t.strip()))
         for t in filter(None, elem.get("ATT", "").split("|")):
@@ -32,7 +34,8 @@ def parse_xml(filepath):
     return nodes, sup_edges, att_edges
 
 
-def parse_json(filepath):
+def parse_json_inline(filepath):
+    """Format A: sections with 'text' fields containing inline <prem>/<conc> XML tags."""
     with open(filepath, "r", encoding="utf-8") as f:
         data = json.load(f)
     nodes, sup_edges, att_edges = {}, [], []
@@ -46,7 +49,7 @@ def parse_json(filepath):
         if not id_m:
             continue
         nid = id_m.group(1)
-        nodes[nid] = tag
+        nodes[nid] = tag.lower()
         sup_m = re.search(r'SUP="([^"]+)"', attrs)
         if sup_m:
             for t in filter(None, sup_m.group(1).split("|")):
@@ -58,13 +61,126 @@ def parse_json(filepath):
     return nodes, sup_edges, att_edges, data.get("document", "")
 
 
+def _clean_id(raw_id):
+    """Convert '[3]' -> '3', strip brackets."""
+    return raw_id.strip().strip("[]").strip()
+
+
+def parse_json_result(filepath):
+    """
+    Format B: pipeline result JSON.
+    Structure: { doc_name: { section_1: { segment1: {id, role, sup, att, text, ...} } } }
+    IDs restart per section, so we prefix them: S1_3, S2_1, etc.
+    """
+    with open(filepath, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    doc_name = list(data.keys())[0]
+    doc      = data[doc_name]
+
+    nodes, sup_edges, att_edges = {}, [], []
+
+    # Section index map: section_1 -> 1, section_2 -> 2, ...
+    sec_index = {}
+    for i, sec_k in enumerate(sorted(doc.keys(), key=lambda k: int(re.search(r'\d+', k).group())), 1):
+        sec_index[sec_k] = i
+
+    # First pass: register all nodes with prefixed IDs
+    # node_map[sec_k][raw_id] = prefixed_id
+    node_map = {}
+    for sec_k, sec in doc.items():
+        si = sec_index[sec_k]
+        node_map[sec_k] = {}
+        for seg in sec.values():
+            raw_id = _clean_id(seg.get("id", ""))
+            if not raw_id:
+                continue
+            role = seg.get("role", "PREM").upper()
+            ntype = "conc" if role == "CONC" else "prem"
+            pid = f"S{si}_{raw_id}"
+            node_map[sec_k][raw_id] = pid
+            # Store extra info for tooltip
+            nodes[pid] = {
+                "type":   ntype,
+                "text":   seg.get("text", "")[:120],
+                "scheme": seg.get("scheme", ""),
+                "atype":  seg.get("type", ""),
+            }
+
+    # Second pass: build edges using prefixed IDs
+    for sec_k, sec in doc.items():
+        si = sec_index[sec_k]
+        nmap = node_map[sec_k]
+        for seg in sec.values():
+            raw_id = _clean_id(seg.get("id", ""))
+            if not raw_id or raw_id not in nmap:
+                continue
+            src_pid = nmap[raw_id]
+
+            sup_raw = seg.get("sup", "")
+            if sup_raw:
+                for t in sup_raw.split(","):
+                    t = _clean_id(t)
+                    if t and t in nmap:
+                        sup_edges.append((src_pid, nmap[t]))
+
+            att_raw = seg.get("att", "")
+            if att_raw:
+                for t in att_raw.split(","):
+                    t = _clean_id(t)
+                    if t and t in nmap:
+                        att_edges.append((src_pid, nmap[t]))
+
+    return nodes, sup_edges, att_edges, doc_name
+
+
+def detect_and_parse(filepath):
+    """Auto-detect JSON format and dispatch to correct parser."""
+    with open(filepath, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    # Format B: single top-level key whose value has section_N keys
+    # containing segment dicts with an 'id' field like '[1]'
+    if isinstance(data, dict) and len(data) == 1:
+        inner = list(data.values())[0]
+        if isinstance(inner, dict):
+            first_val = next(iter(inner.values()), None)
+            if isinstance(first_val, dict):
+                first_seg = next(iter(first_val.values()), None)
+                if isinstance(first_seg, dict) and "id" in first_seg:
+                    nodes, sup_edges, att_edges, doc_name = parse_json_result(filepath)
+                    # nodes here is dict of pid -> {type, text, ...}
+                    # convert to simple type map for graph builder
+                    node_types = {k: v["type"] for k, v in nodes.items()}
+                    node_meta  = nodes
+                    return node_types, sup_edges, att_edges, doc_name, node_meta
+
+    # Format A: inline XML tags in text fields
+    nodes, sup_edges, att_edges, doc_name = parse_json_inline(filepath)
+    node_types = nodes
+    node_meta  = {k: {"type": v, "text": "", "scheme": "", "atype": ""} for k, v in nodes.items()}
+    return node_types, sup_edges, att_edges, doc_name, node_meta
+
+
 # ── HTML builder ──────────────────────────────────────────────────────────────
 
-def build_graph_html(nodes, sup_edges, att_edges, doc_name=""):
-    nodes_js = json.dumps([{"id": k, "type": v} for k, v in nodes.items()])
-    sup_js   = json.dumps([{"source": s, "target": t} for s, t in sup_edges])
-    att_js   = json.dumps([{"source": s, "target": t} for s, t in att_edges])
-    label    = doc_name.replace("_", " ") if doc_name else "argumentation graph"
+def build_graph_html(node_types, sup_edges, att_edges, doc_name="", node_meta=None):
+    if node_meta is None:
+        node_meta = {}
+
+    nodes_js = json.dumps([
+        {
+            "id":     k,
+            "type":   v,
+            "text":   node_meta.get(k, {}).get("text", ""),
+            "scheme": node_meta.get(k, {}).get("scheme", ""),
+            "atype":  node_meta.get(k, {}).get("atype", ""),
+        }
+        for k, v in node_types.items()
+    ])
+    sup_js = json.dumps([{"source": s, "target": t} for s, t in sup_edges])
+    att_js = json.dumps([{"source": s, "target": t} for s, t in att_edges])
+    label  = doc_name.replace("_", " ") if doc_name else "argumentation graph"
 
     return """<!DOCTYPE html>
 <html lang="en">
@@ -117,8 +233,7 @@ def build_graph_html(nodes, sup_edges, att_edges, doc_name=""):
   .link-sup { stroke:var(--green); stroke-width:1.8; stroke-opacity:.75; fill:none; }
   .link-att { stroke:var(--red);   stroke-width:2.2; stroke-opacity:.9;
               stroke-dasharray:7 3; fill:none; }
-  .ng { cursor:grab; }
-  .ng:active { cursor:grabbing; }
+  .ng { cursor:grab; } .ng:active { cursor:grabbing; }
   .cp { fill:var(--blue-bg); stroke:var(--blue-br); stroke-width:1.5; transition:fill .15s; }
   .cc { fill:var(--pur-bg);  stroke:var(--pur-br);  stroke-width:1.5; transition:fill .15s; }
   .ng:hover .cp { fill:var(--blue-br); }
@@ -127,10 +242,10 @@ def build_graph_html(nodes, sup_edges, att_edges, doc_name=""):
         text-anchor:middle; dominant-baseline:central; pointer-events:none; }
   #tt {
     position:fixed; display:none; background:#1c2030;
-    border:1px solid var(--border); border-radius:7px; padding:9px 13px;
-    font-family:'IBM Plex Mono',monospace; font-size:11px; line-height:1.75;
-    pointer-events:none; z-index:999; max-width:260px; color:var(--text);
-    box-shadow:0 8px 32px rgba(0,0,0,.55);
+    border:1px solid var(--border); border-radius:7px; padding:10px 14px;
+    font-family:'IBM Plex Mono',monospace; font-size:11px; line-height:1.8;
+    pointer-events:none; z-index:999; max-width:300px; color:var(--text);
+    box-shadow:0 8px 32px rgba(0,0,0,.6);
   }
 </style>
 </head>
@@ -173,16 +288,15 @@ const g = svg.append("g");
 svg.call(d3.zoom().scaleExtent([0.1, 5]).on("zoom", e => g.attr("transform", e.transform)));
 
 const sim = d3.forceSimulation(NODES)
-  .force("link",      d3.forceLink(links).id(d => d.id).distance(85).strength(0.55))
-  .force("charge",    d3.forceManyBody().strength(-340))
+  .force("link",      d3.forceLink(links).id(d => d.id).distance(90).strength(0.5))
+  .force("charge",    d3.forceManyBody().strength(-360))
   .force("center",    d3.forceCenter(W / 2, H / 2))
-  .force("collision", d3.forceCollide(20));
+  .force("collision", d3.forceCollide(22));
 
 const lsel = g.append("g").selectAll("line").data(links).join("line")
   .attr("class",      d => d.kind === "sup" ? "link-sup" : "link-att")
   .attr("marker-end", d => d.kind === "sup" ? "url(#ms)" : "url(#ma)");
 
-// adjacency for tooltip
 const adj = {};
 NODES.forEach(n => { adj[n.id] = {sf:[], st:[], af:[], at:[]}; });
 SUP.forEach(e => { adj[e.target].sf.push(e.source); adj[e.source].st.push(e.target); });
@@ -198,7 +312,12 @@ const nsel = g.append("g").selectAll("g").data(NODES).join("g")
     .on("end",   (e, d) => { if (!e.active) sim.alphaTarget(0); d.fx = null; d.fy = null; }))
   .on("mouseover", (e, d) => {
     const a = adj[d.id];
-    let h = "<b style='font-size:13px'>" + d.id + "</b> <span style='color:#64748b'>" + d.type + "</span><hr style='border:none;border-top:1px solid #2a2f42;margin:5px 0'>";
+    let h = "<b style='font-size:13px'>" + d.id + "</b>";
+    h += " <span style='color:#64748b'>" + d.type + "</span>";
+    if (d.scheme) h += " <span style='color:#94a3b8'>· " + d.scheme + "</span>";
+    if (d.atype)  h += " <span style='color:#94a3b8'>· type:" + d.atype + "</span>";
+    h += "<hr style='border:none;border-top:1px solid #2a2f42;margin:5px 0'>";
+    if (d.text) h += "<div style='color:#94a3b8;font-size:10px;margin-bottom:5px;line-height:1.5'>" + d.text + (d.text.length >= 120 ? "…" : "") + "</div>";
     if (a.sf.length) h += "<div style='color:#22c55e'>◀ SUP from: " + a.sf.join(", ") + "</div>";
     if (a.st.length) h += "<div style='color:#22c55e'>▶ SUP to: "   + a.st.join(", ") + "</div>";
     if (a.af.length) h += "<div style='color:#ef4444'>◀ ATT from: " + a.af.join(", ") + "</div>";
@@ -211,9 +330,9 @@ const nsel = g.append("g").selectAll("g").data(NODES).join("g")
   .on("mousemove", e => { tt.style.left = (e.clientX + 14) + "px"; tt.style.top = (e.clientY - 12) + "px"; })
   .on("mouseleave", () => { tt.style.display = "none"; });
 
-nsel.append("circle").attr("r", 13).attr("class", d => d.type === "prem" ? "cp" : "cc");
+nsel.append("circle").attr("r", 14).attr("class", d => d.type === "prem" ? "cp" : "cc");
 nsel.append("text").attr("class", "nl")
-  .style("font-size", d => d.id.length > 4 ? "7px" : "9px")
+  .style("font-size", d => d.id.length > 5 ? "6.5px" : "8px")
   .text(d => d.id);
 
 sim.on("tick", () => {
@@ -226,10 +345,7 @@ sim.on("tick", () => {
 </html>"""
 
 
-# ── File handler ──────────────────────────────────────────────────────────────
-
-# We save the HTML to a temp file and serve it via gr.File so it opens in browser,
-# AND embed it in an iframe inside gr.HTML for inline preview.
+# ── Gradio handler ────────────────────────────────────────────────────────────
 
 def process_file(file_obj):
     if file_obj is None:
@@ -241,26 +357,27 @@ def process_file(file_obj):
     try:
         if ext == "xml":
             nodes, sup_edges, att_edges = parse_xml(path)
-            doc_name = os.path.basename(path).replace(".xml", "").replace("_", " ")
+            doc_name  = os.path.basename(path).replace(".xml", "").replace("_", " ")
+            node_meta = {k: {"type": v, "text": "", "scheme": "", "atype": ""} for k, v in nodes.items()}
+            node_types = nodes
         elif ext == "json":
-            nodes, sup_edges, att_edges, doc_name = parse_json(path)
+            node_types, sup_edges, att_edges, doc_name, node_meta = detect_and_parse(path)
         else:
             return None, "<p style='color:#ef4444;font-family:monospace;padding:20px'>Upload .xml or .json</p>"
     except Exception as ex:
-        return None, f"<p style='color:#ef4444;font-family:monospace;padding:20px'>Parse error: {ex}</p>"
+        import traceback
+        return None, f"<pre style='color:#ef4444;font-family:monospace;padding:20px'>{traceback.format_exc()}</pre>"
 
-    if not nodes:
-        return None, "<p style='color:#ef4444;font-family:monospace;padding:20px'>No nodes found.</p>"
+    if not node_types:
+        return None, "<p style='color:#ef4444;font-family:monospace;padding:20px'>No nodes found in file.</p>"
 
-    html = build_graph_html(nodes, sup_edges, att_edges, doc_name)
+    html = build_graph_html(node_types, sup_edges, att_edges, doc_name, node_meta)
 
-    # Write to a named temp file
     tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".html", mode="w", encoding="utf-8")
     tmp.write(html)
     tmp.close()
 
-    # Inline iframe preview (srcdoc)
-    srcdoc = html.replace('"', "&quot;")
+    srcdoc = html.replace("&", "&amp;").replace('"', "&quot;")
     iframe = f'<iframe srcdoc="{srcdoc}" style="width:100%;height:680px;border:none;border-radius:8px;"></iframe>'
 
     return tmp.name, iframe
@@ -294,7 +411,7 @@ with gr.Blocks(
     with gr.Row(elem_id="title-row"):
         gr.HTML("""
         <h1>⚖ Argumentation Graph Visualizer</h1>
-        <p>Upload a Demosthenes <b>.xml</b> or structured <b>.json</b> file
+        <p>Upload a Demosthenes <b>.xml</b> or pipeline result <b>.json</b>
            → interactive force-directed argument graph.</p>
         """)
 
@@ -302,7 +419,7 @@ with gr.Blocks(
         with gr.Column(scale=1, min_width=260):
             file_input = gr.File(label="Upload XML or JSON", file_types=[".xml", ".json"])
             run_btn    = gr.Button("Build Graph", variant="primary", size="lg")
-            dl_out     = gr.File(label="Download graph HTML", visible=True)
+            dl_out     = gr.File(label="Download graph HTML")
             gr.HTML("""
             <div style="font-family:monospace;font-size:11px;color:#64748b;
                         border:1px solid #2a2f42;border-radius:8px;padding:12px;margin-top:6px;">
@@ -330,8 +447,9 @@ with gr.Blocks(
                   <circle cx="6" cy="6" r="5" fill="#4c1d95" stroke="#a855f7" stroke-width="1.5"/>
                 </svg> conc (conclusion)
               </div>
-              <div style="margin-top:8px;color:#475569;">
-                Drag · Scroll to zoom · Hover node
+              <div style="margin-top:8px;color:#475569;line-height:1.6;">
+                Accepts .xml · .json (inline tags) · .json (pipeline result)<br><br>
+                Hover node �� see text snippet, type, scheme, relations
               </div>
             </div>
             """)