qare_engine.py

"""
qare_engine.py — the QARE runtime ("firmware").

WHY THIS EXISTS
---------------
The fine-tuned adapter was trained on prompts that CONTAIN a precomputed
"Computed resource requirements" block. If you hand the model a bare question,
it will hallucinate requirements (measured: it claimed Grover needs ~2057 qubits
for a 20-bit search; the true value is 21). So the model is NOT standalone.

QARE is therefore a HYBRID system:
    knowledge_base.py  -> deterministic resource estimation + feasibility  (the decision)
    BitNet + LoRA      -> natural-language ranking, reasoning, explanation  (the narration)

This module owns that contract: it computes the facts, builds the exact prompt
format the model was trained on, generates, and returns both the raw text and a
parsed result. Anything that skips this and prompts the model directly is
misusing it.

USAGE
-----
    from qare_engine import QAREEngine
    eng = QAREEngine(base="microsoft/bitnet-b1.58-2B-4T-bf16",
                     adapter="Uluka/qare-bitnet-lora")
    out = eng.recommend(problem_type="unstructured_search", size=20,
                        available_qubits=1000, noise="none", max_depth=100000,
                        hardware="simulator", desired_accuracy=0.9)
    print(out["text"])

CLI
---
    python qare_engine.py --problem unstructured_search --size 20 --qubits 1000 \
        --noise none --depth 100000 --hardware simulator --accuracy 0.9
"""
from __future__ import annotations

import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from knowledge_base import (  # noqa: E402
    ALGOS, CANDIDATES, HARDWARE_TYPES, NOISE_LEVELS, PROBLEM_TYPES,
    Problem, _feasible, recommend,
)

SYSTEM = ("You are QARE, a quantum algorithm recommendation engine. You are given a "
          "computational problem, hardware constraints, and PRECOMPUTED resource "
          "requirements for each candidate algorithm. Select the best algorithm using "
          "ONLY the numbers provided -- do not invent requirements. Recommend a "
          "classical algorithm when no quantum method fits. Respond in the fixed "
          "QARE format, starting with 'Primary Algorithm:'.")

SIZE_WORD = {
    "integer_factoring": "bit integer", "discrete_log": "bit modulus",
    "unstructured_search": "item database (log2)",
    "combinatorial_optimization": "binary variables",
    "ground_state_energy": "spin-orbitals",
    "eigenvalue_estimation": "dimension (log2) operator",
    "linear_system": "dimension (log2) system", "sampling": "modes",
    "classification": "features", "graph_connectivity": "nodes (log2)",
    "simulation_dynamics": "sites",
}


def build_requirements(p: Problem) -> str:
    """Precomputed facts, in the exact format the adapter was trained on."""
    lines = []
    for key in CANDIDATES.get(p.problem_type, ["Classical"]):
        algo = ALGOS[key]
        if key == "Classical":
            lines.append("  Classical algorithm: 0 qubits, no circuit "
                         "-> always available")
            continue
        req_q = algo.qubits_fn(p.size)
        req_d = algo.depth_fn(p.size)
        feas, fails = _feasible(algo, p)
        verdict = "FITS" if feas else "BLOCKED: " + "; ".join(fails)
        lines.append(f"  {algo.name}: needs {req_q} qubits, depth {req_d} -> {verdict}")
    return "\n".join(lines)


def build_prompt(p: Problem) -> str:
    unit = SIZE_WORD.get(p.problem_type, "unit")
    return ("Recommend the best quantum algorithm for this problem.\n\n"
            f"Problem: {p.problem_type.replace('_', ' ')} on a {p.size}-{unit} instance.\n"
            f"Hardware: {p.hardware}\n"
            f"Available qubits: {p.available_qubits}\n"
            f"Noise level: {p.noise}\n"
            f"Max circuit depth: {p.max_depth}\n"
            f"Desired accuracy: {p.desired_accuracy}\n\n"
            "Computed resource requirements (from a resource estimator):\n"
            f"{build_requirements(p)}")


class QAREEngine:
    def __init__(self,
                 base: str = "microsoft/bitnet-b1.58-2B-4T-bf16",
                 adapter: str | None = None,
                 device_map: str = "auto"):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        # Colab/older envs ship torchao 0.10.0; PEFT raises instead of skipping.
        try:
            import peft.import_utils as _iu
            _iu.is_torchao_available = lambda: False
            import peft.tuners.lora.torchao as _lt
            _lt.is_torchao_available = lambda: False
        except Exception:
            pass

        # Tokenizer always from BASE (LoRA never changes it).
        self.tok = AutoTokenizer.from_pretrained(base)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token

        # BitNet REQUIRES bf16: fp16 overflows BitLinear -> NaN.
        self.model = AutoModelForCausalLM.from_pretrained(
            base, torch_dtype=torch.bfloat16, device_map=device_map)

        if adapter:
            from peft import PeftModel
            self.model = PeftModel.from_pretrained(self.model, adapter)
        self.model.eval()
        self.torch = torch

    def recommend(self, problem_type: str, size: int, available_qubits: int,
                  noise: str, max_depth: int, hardware: str,
                  desired_accuracy: float = 0.95,
                  max_new_tokens: int = 320) -> dict:
        if problem_type not in PROBLEM_TYPES:
            raise ValueError(f"problem_type must be one of {PROBLEM_TYPES}")
        if hardware not in HARDWARE_TYPES:
            raise ValueError(f"hardware must be one of {HARDWARE_TYPES}")
        if noise not in NOISE_LEVELS:
            raise ValueError(f"noise must be one of {NOISE_LEVELS}")

        p = Problem(problem_type, size, available_qubits, noise, max_depth,
                    hardware, desired_accuracy)
        user = build_prompt(p)
        msgs = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": user}]
        text = self.tok.apply_chat_template(msgs, add_generation_prompt=True,
                                            tokenize=False)
        enc = self.tok(text, return_tensors="pt").to(self.model.device)
        with self.torch.no_grad():
            out = self.model.generate(
                **enc, max_new_tokens=max_new_tokens, do_sample=False,
                pad_token_id=self.tok.pad_token_id or self.tok.eos_token_id)
        gen = self.tok.decode(out[0][enc["input_ids"].shape[1]:],
                              skip_special_tokens=True)

        return {
            "text": gen,
            "prompt": user,
            # Deterministic KB answer — the audit trail / ground truth.
            "kb_reference": recommend(p),
        }


def main():
    ap = argparse.ArgumentParser(description="QARE runtime")
    ap.add_argument("--base", default="microsoft/bitnet-b1.58-2B-4T-bf16")
    ap.add_argument("--adapter", default=None, help="LoRA adapter path or HF repo id")
    ap.add_argument("--problem", required=True, choices=PROBLEM_TYPES)
    ap.add_argument("--size", type=int, required=True)
    ap.add_argument("--qubits", type=int, required=True)
    ap.add_argument("--noise", default="medium", choices=NOISE_LEVELS)
    ap.add_argument("--depth", type=int, default=500)
    ap.add_argument("--hardware", default="superconducting", choices=HARDWARE_TYPES)
    ap.add_argument("--accuracy", type=float, default=0.95)
    ap.add_argument("--show_prompt", action="store_true")
    ap.add_argument("--show_kb", action="store_true",
                    help="print the deterministic KB answer for comparison")
    a = ap.parse_args()

    eng = QAREEngine(base=a.base, adapter=a.adapter)
    r = eng.recommend(a.problem, a.size, a.qubits, a.noise, a.depth,
                      a.hardware, a.accuracy)
    if a.show_prompt:
        print("=" * 70 + "\nPROMPT\n" + "=" * 70)
        print(r["prompt"])
    print("=" * 70 + "\nQARE RECOMMENDATION\n" + "=" * 70)
    print(r["text"])
    if a.show_kb:
        import json
        print("=" * 70 + "\nKB REFERENCE (deterministic)\n" + "=" * 70)
        print(json.dumps(r["kb_reference"], indent=2))


if __name__ == "__main__":
    main()