benchmark_runner.py

"""
Benchmark Runner
================
Runs queries + text samples through the full ContextSieveEngine pipeline
and prints:

  1. Raw Tokens vs Compressed Tokens -> Compression Ratio (%)
  2. Model latency -> Latency (ms)
  3. Compressed text -> human-readable check (telegraph vs fluent)

Automatically detects and uses a trained NanoPruner model if available,
otherwise falls back to the mock forward pass.

Usage:
    python benchmark_runner.py
"""

from __future__ import annotations

import sys
import textwrap
import time
from datetime import datetime
from pathlib import Path

from context_sieve.inference import ContextSieveEngine

# ---------------------------------------------------------------------------
# Test fixtures - representative conversational data
# ---------------------------------------------------------------------------
SAMPLES = [
    {
        "query": "What was the cause of the financial discrepancy in Q3?",
        "text": (
            "During the quarterly board review held on September 15th, the CFO "
            "presented the consolidated financial statements. Revenue for Q3 "
            "came in at $42.7 million, which was 3% below the forecast of "
            "$44 million. The primary cause of the discrepancy was identified "
            "as a delayed contract renewal with Meridian Corp, worth approximately "
            "$1.8 million. This delay was attributed to ongoing legal negotiations "
            "regarding intellectual property clauses in the master services "
            "agreement. Additionally, the APAC region underperformed expectations "
            "due to currency headwinds, particularly the weakening of the Japanese "
            "yen against the US dollar. The board discussed several mitigation "
            "strategies, including accelerating the Meridian deal closure and "
            "implementing dynamic currency hedging for APAC receivables. The "
            "audit committee noted that there were no material misstatements in "
            "the financial reports and that internal controls were operating "
            "effectively. Furthermore, the marketing department reported that "
            "the Q3 brand awareness campaign reached 12 million impressions "
            "but only converted at a 0.8% rate, below the 1.2% industry average. "
            "The VP of Sales mentioned that the new CRM system rollout caused "
            "temporary disruptions in the pipeline tracking, leading to delayed "
            "follow-ups on 47 qualified leads. Despite these challenges, the "
            "customer retention rate remained strong at 94.2%, and the NPS score "
            "improved from 62 to 67. The CEO emphasized the importance of "
            "maintaining focus on the strategic initiative codenamed 'Project "
            "Atlas', which aims to expand the company's presence in the European "
            "market through a series of targeted acquisitions. The next quarterly "
            "review is scheduled for December 18th."
        ),
    },
    {
        "query": "Did the defendant have an alibi for the night of March 3rd?",
        "text": (
            "The witness testimony from the preliminary hearing on April 2nd "
            "revealed several conflicting accounts. Detective Morrison stated "
            "that the defendant was seen at the Grand Hotel lobby at approximately "
            "9:47 PM on March 3rd, based on CCTV footage recovered from the "
            "hotel's security system. However, the defendant's attorney presented "
            "a signed affidavit from the defendant's brother, claiming they were "
            "together at a family dinner in Rockville until 10:30 PM. Phone "
            "records obtained via subpoena showed the defendant's mobile device "
            "connected to a cell tower in downtown Baltimore at 9:52 PM, which "
            "is inconsistent with the Rockville alibi but consistent with the "
            "Grand Hotel location. The forensic analyst testified that the "
            "defendant's fingerprints were NOT found at the crime scene, and "
            "DNA evidence collected from the doorknob was inconclusive due to "
            "degradation. The prosecution argued that the defendant had motive "
            "based on a $2.3 million life insurance policy taken out six months "
            "prior. The defense countered that the policy was a routine financial "
            "planning measure recommended by the defendant's wealth advisor, "
            "citing email correspondence between the advisor and the defendant "
            "dating back to January. The judge noted the discrepancy between "
            "the cell tower data and the alibi testimony and ordered additional "
            "discovery. The case was continued to May 15th. Meanwhile, the "
            "victim's family issued a public statement expressing frustration "
            "with the pace of the investigation."
        ),
    },
    {
        "query": "What are the key technical specifications of the new chip?",
        "text": (
            "The engineering team presented the final specifications for the "
            "NX-7200 processor at the internal design review. The chip features "
            "a 3nm process node manufactured by TSMC, with 18 billion transistors "
            "packed into a 145mm2 die area. It supports DDR5-6400 memory with "
            "a maximum bandwidth of 102.4 GB/s across four channels. The CPU "
            "complex consists of 12 high-performance cores running at a base "
            "clock of 3.6 GHz with a boost of 5.4 GHz, alongside 4 efficiency "
            "cores at 2.4 GHz. The integrated GPU features 24 compute units "
            "capable of 8.2 TFLOPS of FP32 performance. Power consumption is "
            "rated at a TDP of 125W with dynamic power management that can "
            "reduce idle consumption to 15W. The chip includes hardware-level "
            "support for AES-256 encryption, secure boot with TPM 2.0 "
            "integration, and a dedicated AI accelerator providing 45 TOPS "
            "for inference workloads. Testing showed that the chip achieves "
            "a Cinebench R24 multi-core score of 42,500, which represents a "
            "23% improvement over the previous generation NX-6100. The thermal "
            "design uses a copper-vapor chamber integrated into the package "
            "substrate, maintaining junction temperatures below 95 degrees Celsius under "
            "sustained load. The team also discussed the supply chain timeline, "
            "noting that initial wafer allocation from TSMC is confirmed for "
            "Q2 production with volume ramp expected in Q3. Marketing requested "
            "that the launch be coordinated with the annual developer conference "
            "in October."
        ),
    },
    {
        "query": "What caused the Redis cache eviction spike at 03:00 AM on production?",
        "text": (
            "The post-mortem incident report for the November 12th outage details the sequence "
            "of events leading to service degradation. At 02:45 AM, CPU utilization on the "
            "primary database cluster spiked to 92% due to an unoptimized analytical query "
            "running on the reporting replica. Simultaneously, the ingress traffic from the "
            "EU-West region increased by 15% following a marketing push. However, the critical "
            "failure occurred at 03:00 AM when the Redis caching layer experienced a massive "
            "spike in key evictions, reaching 45,000 evictions per second. The root cause of "
            "this eviction storm was traced to a misconfigured Kubernetes cron job that executed "
            "a bulk cache invalidation script using the 'KEYS *' command instead of 'SCAN', "
            "completely saturating the Redis single thread and blowing out the memory limit. "
            "The system attempted to auto-scale, but the new pods failed to pass readiness "
            "probes because the configuration map was lacking the updated database credentials. "
            "Service was fully restored at 04:15 AM after rolling back the cron job."
        ),
    },
    {
        "query": "Which AWS service contributed most to the budget overrun in October?",
        "text": (
            "Q4 FinOps Cloud Expenditure Review. The total cloud spend for October reached "
            "$142,500, which is an 18% variance over the projected budget of $120,000. "
            "Base compute costs on EC2 remained relatively stable at $45,000, thanks to our "
            "aggressive transition to Graviton instances. Storage costs on S3 increased slightly "
            "to $12,000 due to the retention of new compliance logs. The RDS Postgres cluster "
            "costs dropped by 5% after we downsized the staging environments. The primary "
            "culprit for the budget overrun was a rogue Amazon SageMaker instance. A data "
            "science experiment was left running on an ml.p4d.24xlarge GPU instance for 14 "
            "continuous days without auto-shutdown policies, accumulating over $22,000 in "
            "unexpected charges. Network egress costs also saw a minor bump due to increased "
            "traffic to our CDN, accounting for an extra $2,000. To prevent future occurrences, "
            "the DevOps team has implemented strict AWS Budgets alerts."
        ),
    },
    {
        "query": "What is the rate limit for the /v2/transactions endpoint on the Pro tier?",
        "text": (
            "Welcome to the Payment Gateway API v2 documentation. Authentication is handled via "
            "Bearer tokens in the Authorization header. Please note that v1 endpoints are "
            "officially deprecated and will be sunset on December 31st. All requests must be "
            "made over HTTPS; plain HTTP requests will be immediately rejected with a 403. "
            "We enforce strict rate limiting to ensure platform stability. For the /v1/users "
            "endpoint, all tiers are capped at 50 requests per second. However, for the newly "
            "optimized /v2/transactions endpoint, limits vary by subscription level. The Basic "
            "tier is restricted to 100 requests per minute and does not support bulk webhooks. "
            "The Pro tier allows up to 1,000 requests per minute and includes priority support "
            "queue access. Enterprise tier limits are custom-configured per tenant and typically "
            "start at 5,000 requests per minute. If a rate limit is exceeded, the API will "
            "return a 429 Too Many Requests response containing a 'Retry-After' header."
        ),
    },
    {
        "query": "What was the observed efficacy rate of the PH-42 drug compared to the placebo at week 24?",
        "text": (
            "Clinical Trial Protocol #7784-B: Phase III Double-Blind Study of PH-42 for chronic "
            "hypertension. A total of 1,250 participants were randomized in a 1:1 ratio. Group A "
            "received 10mg of PH-42 daily, while Group B received a starch-based placebo. "
            "Baseline systolic blood pressure averaged 158 mmHg across both cohorts. After 12 "
            "weeks, Group A showed a Mean arterial pressure (MAP) reduction of 14 mmHg. The "
            "primary endpoint was measured at week 24. Results indicated that 78.4% of patients "
            "in Group A achieved the target BP of <130/80, compared to only 14.2% in the "
            "placebo group (p < 0.001). Adverse events were reported in 4% of the active group, "
            "mostly limited to mild dizziness and transient fatigue. The study concluded that "
            "PH-42 is significantly more effective than placebo for sustained BP management. "
            "Long-term follow-up is planned for another 48 months to monitor cardiovascular events."
        ),
    },
    {
        "query": "What is the remaining fuel margin for the satellite after the orbital correction maneuver?",
        "text": (
            "Mission Operations Update: Helios-9 Communications Satellite. The craft reached its "
            "geostationary transfer orbit (GTO) at 22:15 UTC. Telemetry indicates a perigee of "
            "250km and an apogee of 35,786km. Following the third apogee kick motor (AKM) "
            "burn, the inclination was reduced to 0.05 degrees. The primary orbital correction "
            "maneuver (OCM-1) consumed 42.5kg of hydrazine propellant, slightly exceeding the "
            "nominal estimate of 40.2kg due to atmospheric drag in the perigee pass. Despite "
            "this, the remaining fuel onboard is 184.2kg, providing a safety margin of "
            "approximately 12% over the 15-year mission life requirement. Solar array deployment "
            "was successful, with current generation at 4.2kW. The next station-keeping "
            "maneuver is scheduled for January 14th."
        ),
    },
    {
        "query": "What is the maximum annual reimbursement for professional certifications under the new policy?",
        "text": (
            "Employee Handbook Update - Section 4.2: Continuous Education and Professional "
            "Development. The company encourages all full-time staff to pursue career growth. "
            "Under the updated 2026 benefits package, employees are eligible for tuition "
            "assistance and certification coverage. For accredited university courses, the "
            "reimbursement is capped at $5,000 per academic year, subject to a 'B' grade "
            "minimum. For professional industry certifications (e.g., AWS, PMP, CPA), the "
            "maximum annual reimbursement is $2,500. This includes exam fees and up to one "
            "retake if the first attempt is unsuccessful. Requests must be submitted via the "
            "HR portal at least 30 days prior to enrollment. Note that any employee leaving "
            "the firm within 12 months of receiving payment must repay 50% of the funds."
        ),
    },
    {
        "query": "What is the source IP of the blocked SQL injection attempt according to the firewall logs?",
        "text": (
            "Security Information and Event Management (SIEM) Log Export - Oct 14. At 23:42:01, the "
            "WAF (Web Application Firewall) triggered a high-severity alert on the production "
            "ingress controller. The payload contained suspicious SQL syntax: 'OR 1=1' and "
            "'UNION SELECT', targeting the /api/v1/auth/login endpoint. The request originated "
            "from source IP address 192.168.45.102, which is associated with a known proxy botnet "
            "out of the Southeast Asia region. The firewall rule 'SQLI-BLOCK-04' successfully "
            "intercepted the packet, returning a 403 Forbidden. No data leakage was detected. "
            "The system automatically blacklisted the CIDR range 192.168.45.0/24 for 48 hours. "
            "Security team has been notified for further forensic investigation."
        ),
    },
    {
        "query": "What is the return window for items bought during the promotional Black Friday sale?",
        "text": (
            "Store Policy & Terms of Service. Standard returns are accepted within 30 days of "
            "delivery for a full refund to the original payment method. Items must be in their "
            "original packaging with tags attached. However, for items purchased during "
            "promotional events (Black Friday, Cyber Monday, or clearance sales), the return "
            "window is reduced to 14 days and only store credit is issued. Final-sale items "
            "marked with a red 'No Returns' tag are completely ineligible for any form of "
            "restitution. Shipping costs for returns are covered by the customer unless the "
            "item is defective. International orders follow a different policy with a flat 45-day "
            "window but no prepaid labels."
        ),
    },
    {
        "query": "Which supplier has the longest lead time for the SKU-8805 motherboard components?",
        "text": (
            "Supply Chain & Procurement Inventory report for Q4. We are currently tracking "
            "inventory levels for the high-end workstation line. For the SKU-8805 motherboard "
            "assembly, we rely on three key vendors. TechParts Inc (Shenzhen) provides the main "
            "PCB with a lead time of 14 days and 98% quality pass rate. PrecisionSolder (Taiwan) "
            "supplies the capacitor kits with a lead time of 7 days via air freight. However, "
            "the critical microchip component from GlobalSemi (Germany) is experiencing "
            "unprecedented delays, with a current lead time of 65 days due to raw material "
            "shortages in silicon wafers. This bottleneck is holding up 12,000 units of "
            "production. Procurement is actively looking for alternative sourcing in the US to "
            "shorten the cycle to 30 days."
        ),
    },
    {
        "query": "What is the maximum acceptable thawing time for the mRNA lipid nanoparticles?",
        "text": (
            "Standard Operating Procedure (SOP) Document 44-B outlines the strict handling protocols "
            "for the experimental mRNA-based therapeutics in the Level 3 biosafety wet lab. To maintain "
            "the structural integrity of the delicate encapsulation layer, the vials must be transferred "
            "from the -80°C ultra-low temperature freezers to the preparation cleanroom using insulated "
            "liquid nitrogen dewars. Upon entering the Class 100 biological safety cabinet, the technician "
            "must scan the vial's QR code into the LIMS (Laboratory Information Management System) to log "
            "the chain of custody. The critical phase is the temperature acclimation process prior to "
            "centrifuge operations. The maximum acceptable thawing time for the mRNA lipid nanoparticles "
            "is exactly 45 minutes at room temperature (20-25°C). If the vials are left exposed to ambient "
            "temperatures for more than 50 minutes, the batch is considered compromised, must be logged "
            "as a critical deviation, and securely destroyed via thermal incineration. Following a "
            "successful thaw, the solution must be diluted with a proprietary phosphate-buffered saline "
            "(PBS) at a 1:4 ratio, gently inverted six times (never shaken to avoid shearing forces), "
            "and immediately loaded into the microfluidic mixing apparatus for the final formulation stage."
        ),
    },
    {
        "query": "Which thruster caused the orbital insertion anomaly on the communication satellite?",
        "text": (
            "Mission Control post-flight analysis report for Launch Vehicle VX-9, payload deployment phase. "
            "Liftoff occurred nominally at T-0 with a thrust-to-weight ratio of 1.4 under optimal weather "
            "conditions. Max-Q (maximum aerodynamic pressure) was passed without any structural anomalies "
            "reported by the onboard strain gauges. The first stage booster successfully performed its "
            "boost-back burn and landed on the autonomous droneship. The second stage engine cutoff (SECO) "
            "occurred perfectly on schedule at T+8 minutes. However, during the critical geostationary "
            "transfer orbit (GTO) circularization burn, the telemetry suite reported a sudden and unexpected "
            "drop in attitude control pressure. The engineering team initiated an emergency diagnostic "
            "sequence. The orbital insertion anomaly was directly traced to a stuck solenoid valve on the "
            "B-4 hydrazine reaction control thruster. This mechanical failure caused an asymmetrical thrust "
            "profile, forcing the onboard autonomous flight computer to abort the main engine burn prematurely "
            "to prevent an uncontrollable flat spin. Ground crews eventually managed to stabilize the "
            "satellite's rotation using the redundant C-block and D-block thrusters. Unfortunately, this "
            "recovery maneuver consumed 14% of the station-keeping propellant, effectively reducing the "
            "commercial operational lifespan of the satellite from 15 years to approximately 11.5 years."
        ),
    },
    {
        "query": "What was the total financial penalty imposed for the spoofing algorithm incident?",
        "text": (
            "The internal compliance and risk management audit for Q2 reviewed the high-frequency trading (HFT) "
            "activities of the automated market-making desk. A severe regulatory violation of the Dodd-Frank "
            "Act was flagged regarding the recent deployment of the 'Hydra' execution algorithm. The algorithm "
            "was originally designed to provide liquidity and tighten spreads in the dark pools for mid-cap "
            "tech equities. However, an unhandled race condition in the C++ order management system caused "
            "the bot to submit and immediately cancel tens of thousands of bid orders within microseconds. "
            "This created a false impression of market depth and artificial buying pressure. This specific "
            "behavior, legally classified by regulators as 'spoofing', triggered an automated investigation "
            "by the Securities and Exchange Commission (SEC). After a grueling three-month legal battle and "
            "forensic code review, the firm reached a settlement without formally admitting guilt. The total "
            "financial penalty imposed for the spoofing algorithm incident amounted to $8.5 million in direct "
            "regulatory fines, plus a mandatory $2 million contribution to a federal investor protection fund. "
            "Additionally, the Lead Quant who bypassed the staging environment and directly approved the merge "
            "request to production was terminated with cause, and the entire algorithmic trading division was "
            "placed under a mandatory 6-month code-freeze."
        ),
    }
]


class Logger:
    """Tee stdout to both terminal and result.log."""
    def __init__(self, filename="result.log"):
        self.terminal = sys.stdout
        self.log = open(filename, "a", encoding="utf-8")
        self.encoding = self.terminal.encoding if hasattr(self.terminal, 'encoding') else 'utf-8'

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        self.terminal.flush()
        self.log.flush()

    def isatty(self):
        return hasattr(self.terminal, 'isatty') and self.terminal.isatty()

    def fileno(self):
        return self.terminal.fileno()


def _try_load_trained_engine() -> ContextSieveEngine:
    """
    Try to load a trained NanoPruner model.  If not available,
    fall back to the mock engine.
    """
    rl_dir = Path("models/rl_nanopruner")
    if (rl_dir / "model_int8.onnx").exists() or (rl_dir / "model.onnx").exists():
        print("[*] Trained RL NanoPruner model detected - loading...")
        return ContextSieveEngine.from_trained(str(rl_dir))

    model_dir = Path("models/nanopruner")
    onnx_exists = (
        (model_dir / "model_int8.onnx").exists()
        or (model_dir / "model.onnx").exists()
    )

    if onnx_exists:
        print("[*] Trained SL NanoPruner model detected - loading...")
        return ContextSieveEngine.from_trained(str(model_dir))
    else:
        print("[!] No trained model found - using MOCK forward pass.")
        print("    To train: python -m context_sieve.trainer.data_generator")
        print("              python -m context_sieve.trainer.train --full-pipeline")
        print()
        return ContextSieveEngine(
            calibrated_threshold=0.45,
            entropy_threshold=0.92,
            hard_cap_tokens=1200,
        )


def _compute_dilation_for_text(text: str):
    """
    Compute dilation metadata for a text sample using the offline indexer.
    Returns char-level boolean array or None if spacy isn't available.
    """
    try:
        from context_sieve.indexer import ParentChildIndexer
        indexer = ParentChildIndexer()
        return indexer._compute_dilation_candidates(text)
    except Exception as e:
        print(f"  [warn] Dilation metadata unavailable: {e}")
        return None


def run_benchmark() -> None:
    """Execute the benchmark and print results."""
    sys.stdout = Logger("result.log")

    print(f"\n[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]")
    engine = _try_load_trained_engine()

    # Pre-compute dilation metadata for all samples.
    print("\n[*] Pre-computing POS dilation metadata (offline indexer)...")
    dilation_data = {}
    try:
        from context_sieve.indexer import ParentChildIndexer
        indexer = ParentChildIndexer()
        for i, sample in enumerate(SAMPLES):
            dilation_data[i] = indexer._compute_dilation_candidates(sample["text"])
    except Exception as e:
        print(f"  [warn] Dilation metadata unavailable: {e}")
        for i in range(len(SAMPLES)):
            dilation_data[i] = None
    print("    Done.\n")

    print("=" * 80)
    print("  CONTEXT SIEVE v8.5 - BENCHMARK")
    print("=" * 80)

    total_orig = 0
    total_comp = 0
    total_time = 0.0

    for i, sample in enumerate(SAMPLES):
        # Run with dilation metadata.
        result = engine.run(
            sample["query"], sample["text"],
            dilation_candidates=dilation_data.get(i),
        )
        total_orig += result.total_original_tokens
        total_comp += result.total_compressed_tokens
        total_time += result.total_latency_ms

        print(f"\n{'-' * 80}")
        print(f"  SAMPLE {i + 1}")
        print(f"{'-' * 80}")
        print(f"  Query: {sample['query']}")
        print(f"  Chunks: {len(result.chunk_results)}")
        print()

        for j, cr in enumerate(result.chunk_results):
            print(f"  [Chunk {j+1}]")
            print(f"    Original tokens  : {cr.original_tokens}")
            print(f"    Compressed tokens: {cr.compressed_tokens}")
            print(
                f"    Compression      : {cr.compression_ratio * 100:.1f}% removed"
            )
            print(f"    Entropy          : {cr.entropy:.4f}")
            print(f"    Bypassed         : {cr.bypassed}")
            print(f"    Latency          : {cr.latency_ms:.2f} ms")
            print()

        print(f"  TOTALS")
        print(f"    {result.total_original_tokens} -> {result.total_compressed_tokens} tokens")
        print(
            f"    Compression ratio: {result.overall_compression_ratio * 100:.1f}%"
        )
        print(f"    Pipeline latency : {result.total_latency_ms:.2f} ms")
        print()
        print(f"  COMPRESSED TEXT:")
        wrapped = textwrap.fill(result.final_text, width=76, initial_indent="    ",
                                subsequent_indent="    ")
        print(wrapped)

    # --- Aggregate ---
    print(f"\n{'=' * 80}")
    print(f"  AGGREGATE RESULTS ({len(SAMPLES)} samples)")
    print(f"{'=' * 80}")
    print(f"  Total orig tokens     : {total_orig}")
    print(f"  Total compressed      : {total_comp}")
    agg_ratio = 1.0 - (total_comp / max(total_orig, 1))
    print(f"  Overall compression   : {agg_ratio * 100:.1f}%")
    print(f"  Total pipeline latency: {total_time:.2f} ms")
    avg_latency = total_time / max(len(SAMPLES), 1)
    print(f"  Avg per-sample latency: {avg_latency:.2f} ms")
    has_rl_model = (Path("models/rl_nanopruner/model_int8.onnx").exists()
                    or Path("models/rl_nanopruner/model.onnx").exists())
    has_model = (Path("models/nanopruner/model_int8.onnx").exists()
                 or Path("models/nanopruner/model.onnx").exists())
                 
    if has_rl_model:
        model_str = "Trained RL NanoPruner"
    elif has_model:
        model_str = "Trained SL NanoPruner"
    else:
        model_str = "MOCK (random)"
        
    print(f"  Model                 : {model_str}")
    print(f"{'=' * 80}\n")


if __name__ == "__main__":
    run_benchmark()