corpus.py

"""
Synthetic corpus for the Markov retrieval tutorial.

Ten short company descriptions across five entities. Concepts have been
manually curated to keep the demonstration inspectable. The corpus is
deliberately small enough that every chunk and every concept can be read
in one sitting, and is shaped to expose the tail-data failure mode of
vanilla RAG.

Each entry has:
  doc_id    a human-readable identifier
  entity    the company name (used to group chunks into documents)
  text      the chunk text used for embedding
  concepts  manually curated concept list (stand-in for an extraction step)

Notable corpus properties:
  - ITAMCO, AcmeMfg use canonical vocabulary ("additive manufacturing")
  - TailCo uses variant vocabulary ("metal 3D printing", "laser sintering")
    for the same capability, never saying "additive manufacturing"
  - ChicagoCo has the capabilities but is in Illinois, not Indiana
  - FoodCo is a distractor: shares Indiana and ISO 9001 but is off-topic
"""

CORPUS = [
    {
        "doc_id": "itamco_capabilities",
        "entity": "ITAMCO",
        "text": "ITAMCO operates a 5-axis CNC machining facility in Indiana. "
                "The company manufactures precision gears and splines for "
                "defense primes. ITAR registered.",
        "concepts": ["5-axis machining", "CNC", "Indiana", "precision gears",
                     "defense primes", "ITAR"],
    },
    {
        "doc_id": "itamco_processes",
        "entity": "ITAMCO",
        "text": "Additive manufacturing capability includes metal 3D printing "
                "and laser sintering. Workforce development partnership with "
                "Notre Dame.",
        "concepts": ["additive manufacturing", "metal 3D printing",
                     "laser sintering", "Notre Dame"],
    },
    {
        "doc_id": "acme_overview",
        "entity": "AcmeMfg",
        "text": "AcmeMfg is an Indiana manufacturer of machined components "
                "for aerospace and defense applications. Operates 5-axis "
                "machining centers.",
        "concepts": ["Indiana", "machined components", "aerospace",
                     "defense primes", "5-axis machining"],
    },
    {
        "doc_id": "acme_certifications",
        "entity": "AcmeMfg",
        "text": "AcmeMfg holds AS9100 and ISO 9001 certifications. ITAR "
                "registered. Additive manufacturing services available.",
        "concepts": ["AS9100", "ISO 9001", "ITAR", "additive manufacturing"],
    },
    {
        "doc_id": "tailco_about",
        "entity": "TailCo",
        "text": "TailCo specializes in metal 3D printing and laser sintering "
                "of titanium components. Located in Indiana.",
        "concepts": ["metal 3D printing", "laser sintering", "titanium",
                     "Indiana"],
    },
    {
        "doc_id": "tailco_services",
        "entity": "TailCo",
        "text": "Direct metal deposition and powder bed fusion services. "
                "AS9100 certified for aerospace work.",
        "concepts": ["direct metal deposition", "powder bed fusion",
                     "AS9100", "aerospace"],
    },
    {
        "doc_id": "chicago_overview",
        "entity": "ChicagoCo",
        "text": "Illinois precision manufacturer. CNC machining and "
                "5-axis capability. Aerospace components.",
        "concepts": ["Illinois", "CNC", "5-axis machining", "aerospace"],
    },
    {
        "doc_id": "chicago_processes",
        "entity": "ChicagoCo",
        "text": "Investment casting and EDM. ISO 9001 quality system.",
        "concepts": ["investment casting", "EDM", "ISO 9001"],
    },
    {
        "doc_id": "foodco_about",
        "entity": "FoodCo",
        "text": "FoodCo produces frozen desserts in Indiana. ISO 9001 "
                "food safety certification.",
        "concepts": ["Indiana", "frozen desserts", "ISO 9001"],
    },
    {
        "doc_id": "foodco_distribution",
        "entity": "FoodCo",
        "text": "Regional distribution across the Midwest. Cold chain "
                "logistics.",
        "concepts": ["Midwest", "cold chain logistics"],
    },
    # PrecisionWorks: canonical vocabulary, out of state. Tests whether
    # the chain can prefer in-state entities for location-anchored queries
    # even when out-of-state entities have strong capability vocabulary.
    {
        "doc_id": "precision_capabilities",
        "entity": "PrecisionWorks",
        "text": "PrecisionWorks is an Ohio-based precision machining "
                "shop. CNC and 5-axis machining services for defense "
                "primes. ITAR registered.",
        "concepts": ["Ohio", "CNC", "5-axis machining", "defense primes",
                     "ITAR"],
    },
    {
        "doc_id": "precision_certifications",
        "entity": "PrecisionWorks",
        "text": "AS9100 and ISO 9001 certified. Specializes in precision "
                "gears and complex machined components.",
        "concepts": ["AS9100", "ISO 9001", "precision gears",
                     "machined components"],
    },
    # AeroForge: canonical vocabulary, in state, complementary capability.
    # Strengthens the central cluster.
    {
        "doc_id": "aeroforge_about",
        "entity": "AeroForge",
        "text": "AeroForge operates an Indiana aerospace forging facility. "
                "AS9100 certified, ITAR registered. Supplies defense primes.",
        "concepts": ["Indiana", "aerospace forging", "AS9100", "ITAR",
                     "defense primes"],
    },
    {
        "doc_id": "aeroforge_processes",
        "entity": "AeroForge",
        "text": "Closed-die forging and ring rolling. ISO 9001 quality "
                "system. Aerospace and defense components.",
        "concepts": ["closed-die forging", "ring rolling", "ISO 9001",
                     "aerospace", "defense primes"],
    },
    # PowderTech: vocabulary tail, more extreme than TailCo. Never says
    # "additive manufacturing" or "metal 3D printing"; uses only the
    # specialist process names. Tests whether the chain can bridge the
    # weaker similarity from "additive manufacturing" to "powder bed fusion".
    {
        "doc_id": "powdertech_about",
        "entity": "PowderTech",
        "text": "PowderTech provides powder bed fusion and direct metal "
                "deposition services in Indiana. Specializes in titanium "
                "and Inconel.",
        "concepts": ["powder bed fusion", "direct metal deposition",
                     "Indiana", "titanium", "Inconel"],
    },
    {
        "doc_id": "powdertech_quality",
        "entity": "PowderTech",
        "text": "AS9100 certified for aerospace applications. ITAR "
                "registered for defense work.",
        "concepts": ["AS9100", "aerospace", "ITAR", "defense primes"],
    },
    # LegacyShop: acronym tail. Describes capabilities entirely in defense
    # procurement vocabulary, no canonical capability terms. Tests whether
    # the chain can find this from a jargon query, and (the inverse case)
    # whether canonical entities surface from a jargon query.
    {
        "doc_id": "legacy_about",
        "entity": "LegacyShop",
        "text": "LegacyShop is a DLA-qualified supplier in Indiana. "
                "Manufactures NSN-coded parts to MIL-spec requirements. "
                "FAR and DFAR compliant.",
        "concepts": ["DLA", "Indiana", "NSN-coded parts", "MIL-spec",
                     "FAR/DFAR compliant"],
    },
    {
        "doc_id": "legacy_capabilities",
        "entity": "LegacyShop",
        "text": "Decades of experience supplying the Defense Logistics "
                "Agency. ITAR registered. AS9100 quality system.",
        "concepts": ["Defense Logistics Agency", "ITAR", "AS9100"],
    },
    # MidwestPrecision: third major canonical defense entity. Justified by
    # the statistical thinness of the canonical cluster - any real
    # industrial-base corpus would have many canonical defense suppliers,
    # not just two or three.
    {
        "doc_id": "midwest_overview",
        "entity": "MidwestPrecision",
        "text": "MidwestPrecision is a Michigan-based precision machining "
                "shop. CNC and 5-axis machining for defense primes and "
                "aerospace OEMs. ITAR registered.",
        "concepts": ["Michigan", "CNC", "5-axis machining", "defense primes",
                     "aerospace", "ITAR"],
    },
    {
        "doc_id": "midwest_quality",
        "entity": "MidwestPrecision",
        "text": "AS9100 and ISO 9001 certified. Long-term partnerships with "
                "tier-one defense contractors. Machined components and "
                "precision gears.",
        "concepts": ["AS9100", "ISO 9001", "defense primes",
                     "machined components", "precision gears"],
    },
    # HoosierMetals: a bilingual entity that spans canonical defense
    # vocabulary and DLA procurement vocabulary. Justified by the fact that
    # real companies often serve both prime contractors and the DLA, and
    # our corpus has none. The DLA mention is alongside canonical
    # vocabulary, not engineered as a bridge.
    {
        "doc_id": "hoosier_overview",
        "entity": "HoosierMetals",
        "text": "HoosierMetals is an Indiana CNC machining shop supplying "
                "machined components to defense primes and DLA-qualified "
                "contracts. AS9100 and ITAR registered.",
        "concepts": ["Indiana", "CNC", "machined components",
                     "defense primes", "DLA", "AS9100", "ITAR"],
    },
    {
        "doc_id": "hoosier_capabilities",
        "entity": "HoosierMetals",
        "text": "Specializes in MIL-spec components and 5-axis machining. "
                "Supports both new programs and legacy sustainment work.",
        "concepts": ["MIL-spec", "5-axis machining", "machined components"],
    },
    # Sintron: AM specialist that uses canonical AM vocabulary in detail.
    # Justified by the imbalance in the current corpus, where canonical AM
    # terminology appears only as one-line mentions in ITAMCO and AcmeMfg
    # while specialist AM vocabulary dominates TailCo and PowderTech.
    {
        "doc_id": "sintron_about",
        "entity": "Sintron",
        "text": "Sintron provides industrial additive manufacturing "
                "services including selective laser melting and electron "
                "beam melting. Indiana facility, AS9100 certified.",
        "concepts": ["Indiana", "additive manufacturing",
                     "selective laser melting", "electron beam melting",
                     "AS9100"],
    },
    {
        "doc_id": "sintron_applications",
        "entity": "Sintron",
        "text": "Aerospace and medical applications. Titanium and stainless "
                "steel printing. ISO 9001 quality system.",
        "concepts": ["aerospace", "medical", "titanium", "stainless steel",
                     "ISO 9001"],
    },
]

# Queries used in the tutorial. The first two test cases where vanilla and
# the chain should both perform reasonably; the third and fourth are
# tail-test queries where the chain is expected to recover items that
# vanilla similarity misses or ranks low.
QUERIES = [
    "Indiana additive manufacturing companies",
    "5-axis machining for defense",
    "metal 3D printing in Indiana",
    "titanium aerospace components",
    "DLA suppliers with defense procurement experience",
]

# Concept pairs whose embedding similarity we print as a sanity check
# before running retrieval. Mix of pairs we expect to be similar (3D
# printing variants, geographic neighbors), pairs we expect to be weak
# (acronym to expansion), and a control pair we expect to be far apart.
DIAGNOSTIC_CONCEPT_PAIRS = [
    ("additive manufacturing", "metal 3D printing"),
    ("additive manufacturing", "laser sintering"),
    ("additive manufacturing", "powder bed fusion"),
    ("additive manufacturing", "direct metal deposition"),
    ("additive manufacturing", "selective laser melting"),
    ("additive manufacturing", "electron beam melting"),
    ("metal 3D printing", "powder bed fusion"),
    ("Indiana", "Illinois"),
    ("Indiana", "Ohio"),
    ("Indiana", "Michigan"),
    ("Indiana", "Midwest"),
    ("CNC", "5-axis machining"),
    ("ITAR", "AS9100"),
    ("DLA", "Defense Logistics Agency"),
    ("DLA", "defense primes"),
    ("MIL-spec", "defense primes"),
    ("NSN-coded parts", "machined components"),
    ("frozen desserts", "cold chain logistics"),
    ("frozen desserts", "additive manufacturing"),
]