run_tutorial.py

"""
Run the tutorial: load corpus, build graph and transition matrix, run the
test queries, print transcripts and path decompositions.

Usage:
    python code/run_tutorial.py

The transcript output should match the previous prototype within
floating-point variation.
"""

import sys
import os
import numpy as np

# Allow running as `python code/run_tutorial.py` from the repo root.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from sentence_transformers import SentenceTransformer

from corpus import CORPUS, QUERIES, DIAGNOSTIC_CONCEPT_PAIRS
from graph import MODEL_NAME, embed_corpus, build_graph
from markov import (
    build_transition_matrix,
    build_pi0_from_query,
    rwr,
    path_decomposition,
)
from figures import make_all_figures


# --------------------------------------------------------------------------
# Reporting helpers
# --------------------------------------------------------------------------

def report_top_chunks(g, scores, k=5, label="chunks"):
    n_chunks = g["n_chunks"]
    chunk_scores = scores[:n_chunks]
    order = np.argsort(-chunk_scores)
    print(f"\n  Top {k} {label} by stationary mass:")
    for rank, idx in enumerate(order[:k], 1):
        e = g["corpus"][idx]
        text_preview = e["text"][:80] + ("..." if len(e["text"]) > 80 else "")
        print(f"    {rank}. [{e['entity']:10s}] score={chunk_scores[idx]:.4f}  "
              f"{e['doc_id']}")
        print(f"       \"{text_preview}\"")


def report_top_concepts(g, scores, k=6):
    n_chunks = g["n_chunks"]
    n_docs = g["n_docs"]
    concept_scores = scores[n_chunks + n_docs:]
    order = np.argsort(-concept_scores)
    print(f"\n  Top {k} concepts by stationary mass:")
    for rank, idx in enumerate(order[:k], 1):
        print(f"    {rank}. {g['concepts'][idx]:30s} "
              f"score={concept_scores[idx]:.4f}")


def report_top_documents(g, scores, k=5):
    n_chunks = g["n_chunks"]
    n_docs = g["n_docs"]
    doc_scores = scores[n_chunks:n_chunks + n_docs]
    order = np.argsort(-doc_scores)
    print(f"\n  Top {k} documents (entities) by stationary mass:")
    for rank, idx in enumerate(order[:k], 1):
        print(f"    {rank}. {g['docs'][idx]:15s} score={doc_scores[idx]:.4f}")


def report_path_decomposition(g, decomp, target_chunk_idx, max_k=8):
    e = g["corpus"][target_chunk_idx]
    print(f"\n  Path decomposition for chunk {target_chunk_idx} "
          f"({e['entity']} / {e['doc_id']}):")
    total = 0.0
    for k in range(max_k + 1):
        contrib = decomp[k, target_chunk_idx]
        total += contrib
        bar = "#" * int(contrib * 300)
        print(f"    k={k}:  {contrib:.5f}   ({bar})")
    print(f"    total: {total:.5f}")


def report_vanilla_baseline(model, g, query_text, k=5):
    """Plain cosine query-chunk ranking, as a baseline for comparison."""
    qemb = model.encode([query_text], normalize_embeddings=True)[0]
    sims = g["chunk_emb"] @ qemb
    order = np.argsort(-sims)
    print(f"\n  Vanilla baseline (cosine query-chunk), top {k}:")
    for rank, idx in enumerate(order[:k], 1):
        e = g["corpus"][idx]
        print(f"    {rank}. [{e['entity']:10s}] sim={sims[idx]:.4f}  "
              f"{e['doc_id']}")


def print_diagnostic_concept_similarities(g):
    print("\nSelected concept-concept similarities:")
    for a, b in DIAGNOSTIC_CONCEPT_PAIRS:
        if a in g["concept_to_idx"] and b in g["concept_to_idx"]:
            ai = g["concept_to_idx"][a]
            bi = g["concept_to_idx"][b]
            s = g["concept_concept_sim_mat"][ai, bi]
            print(f"  {a:30s} <-> {b:30s} sim={s:.3f}")


# --------------------------------------------------------------------------
# Main
# --------------------------------------------------------------------------

def main():
    print(f"Loading {MODEL_NAME}...")
    model = SentenceTransformer(MODEL_NAME)

    print(f"  embedding {len(CORPUS)} chunks and "
          f"{len(set(c for entry in CORPUS for c in entry['concepts']))} "
          f"concepts with {MODEL_NAME}...")
    chunk_emb, concept_emb, concept_list = embed_corpus(model, CORPUS)
    g = build_graph(CORPUS, chunk_emb, concept_emb, concept_list)

    print(f"\nGraph: {g['n_chunks']} chunks, {g['n_docs']} entities, "
          f"{g['n_concepts']} concepts, {g['n_total']} total nodes")

    P = build_transition_matrix(g)
    density = 100 * P.nnz / (P.shape[0] * P.shape[1])
    print(f"Transition matrix: {P.shape[0]}x{P.shape[1]}, "
          f"{P.nnz} nonzero ({density:.1f}% dense)")

    print_diagnostic_concept_similarities(g)

    # Capture (pi0, decomposition) per query so figures can use them.
    decomps_by_query = {}

    for qtext in QUERIES:
        print("\n" + "=" * 72)
        print(f'Query: "{qtext}"')
        print("=" * 72)

        report_vanilla_baseline(model, g, qtext, k=5)

        pi0 = build_pi0_from_query(model, g, qtext)
        pi_star, iters = rwr(P, pi0, c=0.15)
        print(f"\n  RWR converged in {iters} iterations")

        report_top_chunks(g, pi_star, k=5)
        report_top_documents(g, pi_star, k=5)
        report_top_concepts(g, pi_star, k=6)

        # Path decomposition for tail-test queries. We trace mass arrival
        # for whichever entities are most informative for each query type.
        decomp_targets = []
        if "metal 3D printing" in qtext or "titanium" in qtext:
            # Vocabulary-tail recovery test: trace TailCo and PowderTech.
            decomp_targets = ["TailCo", "PowderTech"]
        elif "additive manufacturing" in qtext:
            # The vocabulary-tail entities that don't say "additive
            # manufacturing" directly. Tests whether the chain bridges
            # canonical query vocabulary to specialist process vocabulary.
            decomp_targets = ["TailCo", "PowderTech"]
        elif "DLA" in qtext or "defense procurement" in qtext:
            # Acronym-tail recovery test: trace LegacyShop, plus a
            # canonical entity to see if the chain bridges the other way.
            decomp_targets = ["LegacyShop", "ITAMCO"]

        if decomp_targets:
            decomp = path_decomposition(P, pi0, c=0.15, max_k=8)
            decomps_by_query[qtext] = (pi0, decomp)
            for entity_name in decomp_targets:
                target_chunks = [i for i, e in enumerate(g["corpus"])
                                 if e["entity"] == entity_name]
                for ci in target_chunks:
                    report_path_decomposition(g, decomp, ci, max_k=8)

    # Generate figures
    print("\n" + "=" * 72)
    print("Generating figures")
    print("=" * 72)
    figures_dir = os.path.join(
        os.path.dirname(os.path.abspath(__file__)),
        "figures",
    )
    make_all_figures(g, P, model, figures_dir, decomps_by_query)


if __name__ == "__main__":
    main()