Spaces:

rmagesh
/

test_MR

Sleeping

Update app.py (#2)
8519c51 2 months ago
26.2 kB
import os
import sys
from pathlib import Path

import gradio as gr
import spaces
import torch
import torch.nn.functional as F

from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

# ============================================================
# CONFIG
# ============================================================

RSNN_REPO = "saifulsajol/Lattice_Spike"

RSNN_CHECKPOINT = (
    "model_weights/our_model/"
    "latticespike_v5.3_causal_rank1_seed42.pt"
)

LOCAL_RSNN_DIR = "/tmp/Lattice_Spike"

LLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct"

LLM_REVISION = (
    "9213176726f574b556790deb65791e0c5aa438b6"
)

# Exact system prompt used by the training/evaluation pipeline
SYSTEM_PROMPT = "Answer the question concisely and factually."

PROMPT_MAX_LENGTH = 2048


# ============================================================
# HF TOKEN
# ============================================================

HF_TOKEN = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    raise RuntimeError(
        "HF_TOKEN is missing. Add HF_TOKEN under "
        "Hugging Face Space Settings -> Secrets."
    )

# ============================================================
# DOWNLOAD LATTICESPIKE REPOSITORY
# ============================================================

print("Downloading LatticeSpike repository...")

snapshot_download(
    repo_id=RSNN_REPO,
    local_dir=LOCAL_RSNN_DIR,
    allow_patterns=[
        "__init__.py",
        "boxes.py",
        "detector.py",
        "llm_signals.py",
        "losses.py",
        "lstm_dyn.py",
        "metrics.py",
        "rsnn.py",
        "signals.py",
        "surrogate.py",

        # Needed because metrics.py imports evaluation.metrics
        "evaluation/*",

        RSNN_CHECKPOINT,
    ],
)

# Needed for:
#   import Lattice_Spike
# and
#   import evaluation
sys.path.insert(
    0,
    str(Path(LOCAL_RSNN_DIR).parent)
)

sys.path.insert(
    0,
    LOCAL_RSNN_DIR
)


from Lattice_Spike import LatticeSpike


# ============================================================
# LOAD RSNN CHECKPOINT
# ============================================================

CHECKPOINT_PATH = os.path.join(
    LOCAL_RSNN_DIR,
    RSNN_CHECKPOINT
)

print("Loading RSNN checkpoint...")

checkpoint = torch.load(
    CHECKPOINT_PATH,
    map_location="cpu",
    weights_only=False,
)

ckpt_args = checkpoint["args"]

print("RSNN checkpoint loaded.")

print(
    "Checkpoint:",
    "d_model =", checkpoint["d_model"],
    "layers =", checkpoint["num_layers"],
    "vocab =", checkpoint["vocab_size"],
)

# ============================================================
# LOAD EXACT LLAMA MODEL USED DURING TRAINING
# ============================================================

print("Loading Llama-3.2-1B-Instruct...")

tok = AutoTokenizer.from_pretrained(
    LLM_MODEL,
    revision=LLM_REVISION,
    token=HF_TOKEN,
)

if tok.pad_token_id is None:
    tok.pad_token_id = tok.eos_token_id


lm = AutoModelForCausalLM.from_pretrained(
    LLM_MODEL,
    revision=LLM_REVISION,
    token=HF_TOKEN,
    torch_dtype=torch.bfloat16,
    attn_implementation="eager",
    low_cpu_mem_usage=True,
)

lm.eval()

# ZeroGPU CUDA emulation works at module scope
lm.to("cuda")

# ============================================================
# VERIFY LLM MATCHES CHECKPOINT
# ============================================================

if int(lm.config.hidden_size) != int(
    checkpoint["d_model"]
):
    raise RuntimeError(
        f"LLM hidden_size={lm.config.hidden_size}, "
        f"checkpoint expects {checkpoint['d_model']}"
    )


if int(lm.config.num_hidden_layers) != int(
    checkpoint["num_layers"]
):
    raise RuntimeError(
        f"LLM layers={lm.config.num_hidden_layers}, "
        f"checkpoint expects {checkpoint['num_layers']}"
    )


if int(lm.config.vocab_size) != int(
    checkpoint["vocab_size"]
):
    raise RuntimeError(
        f"LLM vocab={lm.config.vocab_size}, "
        f"checkpoint expects {checkpoint['vocab_size']}"
    )


print("Llama architecture matches checkpoint.")


# ============================================================
# BUILD LATTICESPIKE
# ============================================================

print("Building LatticeSpike...")


rsnn = LatticeSpike(

    d_model=int(
        checkpoint["d_model"]
    ),

    vocab_size=int(
        checkpoint["vocab_size"]
    ),

    d_box=int(
        ckpt_args["d_box"]
    ),

    num_layers=int(
        checkpoint["num_layers"]
    ),

    mrl_k=int(
        ckpt_args["mrl_k"]
    ),

    n_neurons=int(
        ckpt_args["n_neurons"]
    ),

    window=int(
        ckpt_args["window"]
    ),

    anomaly_window=int(
        ckpt_args["anomaly_window"]
    ),

    rho_star=float(
        ckpt_args["rho_star"]
    ),

    surrogate=ckpt_args[
        "surrogate"
    ],

    concretization_scale=float(
        ckpt_args["concretization_scale"]
    ),

    # Checkpoint uses Channel 6
    include_concretization_channel=True,

    tie_box_token_emb=bool(
        ckpt_args["tie_box_token_emb"]
    ),

    # IMPORTANT:
    # use the real Llama embedding matrix
    token_embedding_table=(
        lm.get_input_embeddings().weight
    ),

    box_encoder=ckpt_args[
        "box_encoder"
    ],

    ablation=ckpt_args[
        "ablation"
    ],

    lstm_use_volume=bool(
        ckpt_args["lstm_use_volume"]
    ),

    lstm_hidden=ckpt_args[
        "lstm_hidden"
    ],

    # This checkpoint has no LayerMixer parameters
    box_hidden_layers=1,

    # No Lookback channel in this checkpoint
    lookback_dim=0,

    # Checkpoint has no learned score_mlp
    score_head="residual",
)


rsnn.load_state_dict(
    checkpoint["model"],
    strict=True
)


# Restore validation-selected threshold
selected_threshold = float(
    checkpoint[
        "threshold_selection"
    ]["selected"]["threshold"]
)

rsnn.anomaly_star.fill_(
    selected_threshold
)

rsnn.eval()
rsnn.to("cuda")


print("RSNN READY")
print(
    "anomaly_star:",
    float(rsnn.anomaly_star)
)
print(
    "rho_star:",
    float(rsnn.rho_star)
)


# ============================================================
# EFFECTIVE RANK
# ============================================================

def effective_rank(H):

    s = torch.linalg.svdvals(
        H.float()
    )

    p = s / s.sum().clamp_min(
        1e-9
    )

    entropy = -(
        p *
        p.clamp_min(1e-9).log()
    ).sum()

    return torch.exp(entropy)


# ============================================================
# BUILD EXACT LLAMA CHAT PROMPT
# ============================================================

def build_prompt(message):

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": message,
        },
    ]

    prompt_text = tok.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    encoded = tok(
        prompt_text,
        add_special_tokens=True,
        truncation=True,
        max_length=PROMPT_MAX_LENGTH,
        return_tensors="pt",
    )

    return {

        "input_ids":
            encoded[
                "input_ids"
            ].to("cuda"),

        "attention_mask":
            encoded[
                "attention_mask"
            ].to("cuda"),
    }


# ============================================================
# REMOVE TERMINAL SPECIAL TOKEN
# ============================================================

def strip_terminal_tokens(answer_ids):

    stop_ids = set()

    if tok.eos_token_id is not None:

        if isinstance(
            tok.eos_token_id,
            (list, tuple)
        ):
            stop_ids.update(
                int(x)
                for x
                in tok.eos_token_id
            )
        else:
            stop_ids.add(
                int(tok.eos_token_id)
            )

    try:

        eot = tok.convert_tokens_to_ids(
            "<|eot_id|>"
        )

        if (
            isinstance(eot, int)
            and eot >= 0
        ):
            stop_ids.add(eot)

    except Exception:
        pass


    while (
        answer_ids.shape[1] > 0
        and
        int(answer_ids[0, -1])
        in stop_ids
    ):

        answer_ids = answer_ids[
            :, :-1
        ]

    return answer_ids


# ============================================================
# CAUSAL SIGNAL EXTRACTION
#
# This matches the training protocol:
# state predicting token t is used to score token t.
# ============================================================

@torch.no_grad()
def extract_causal_signals(
    prompt_ids,
    answer_ids,
):

    if answer_ids.shape[1] == 0:

        raise ValueError(
            "No answer tokens generated."
        )


    input_ids = torch.cat(
        [
            prompt_ids,
            answer_ids,
        ],
        dim=1,
    )


    outputs = lm(

        input_ids=input_ids,

        output_hidden_states=True,

        output_attentions=True,

        use_cache=False,
    )


    hidden = outputs.hidden_states

    if hidden is None:

        raise RuntimeError(
            "LLM did not return hidden states."
        )


    num_layers = (
        len(hidden) - 1
    )

    prompt_len = int(
        prompt_ids.shape[1]
    )

    answer_len = int(
        answer_ids.shape[1]
    )


    # ========================================================
    # CAUSAL ALIGNMENT
    #
    # prompt_len - 1 predicts first answer token
    # ========================================================

    pred_positions = torch.arange(

        prompt_len - 1,

        prompt_len
        + answer_len
        - 1,

        device=input_ids.device,

        dtype=torch.long,
    )


    # Checkpoint default:
    # middle Llama layer
    layer_idx = max(
        1,
        num_layers // 2
    )


    h_box = (
        hidden[layer_idx][0]
        .index_select(
            0,
            pred_positions
        )
        .float()
    )


    # ========================================================
    # CHANNEL 3:
    # Per-layer KL trace
    # ========================================================

    lm_head = (
        lm.get_output_embeddings()
    )


    delta = torch.zeros(

        answer_len,

        num_layers,

        device=input_ids.device,

        dtype=torch.float32,
    )


    for layer_no in range(
        1,
        num_layers + 1
    ):

        layer_hidden = (
            hidden[layer_no][0]
            .index_select(
                0,
                pred_positions
            )
        )


        logits = (
            lm_head(
                layer_hidden
            )
            .float()
        )


        logp = F.log_softmax(
            logits,
            dim=-1,
        )


        probs = logp.exp()


        if answer_len > 1:

            delta[
                1:,
                layer_no - 1
            ] = (

                probs[1:]
                *
                (
                    logp[1:]
                    -
                    logp[:-1]
                )

            ).sum(
                dim=-1
            )


        del (
            layer_hidden,
            logits,
            logp,
            probs,
        )


    # ========================================================
    # ATTENTION ENTROPY
    # ========================================================

    attn_ent = torch.zeros(

        answer_len,

        device=input_ids.device,

        dtype=torch.float32,
    )


    valid_attns = [

        attn

        for attn
        in (
            outputs.attentions
            or ()
        )

        if attn is not None
    ]


    for attn in valid_attns:

        answer_attn = (

            attn[0]

            .index_select(
                1,
                pred_positions
            )

            .float()
        )


        entropy = -(

            answer_attn

            *
            answer_attn
            .clamp_min(
                1e-9
            )
            .log()

        ).sum(
            dim=-1
        )


        attn_ent += (
            entropy.mean(
                dim=0
            )
        )


    attn_ent /= max(
        len(valid_attns),
        1
    )


    # ========================================================
    # EFFECTIVE RANK
    # ========================================================

    layer_stack = torch.stack(

        [

            hidden[layer_no][0]

            .index_select(
                0,
                pred_positions
            )

            .float()

            for layer_no
            in range(
                1,
                num_layers + 1
            )
        ],

        dim=0,
    )


    erank = torch.zeros(

        answer_len,

        device=input_ids.device,

        dtype=torch.float32,
    )


    for t in range(
        answer_len
    ):

        erank[t] = effective_rank(

            layer_stack[
                :, t, :
            ]
        )


    # ========================================================
    # TOP-K NEXT TOKEN DISTRIBUTION
    # ========================================================

    pred_hidden = (

        hidden[-1][0]

        .index_select(
            0,
            pred_positions
        )
    )


    final_logits = (

        lm_head(
            pred_hidden
        )

        .float()
    )


    top_k = int(
        ckpt_args["top_k"]
    )


    topk_probs, topk_ids = (

        F.softmax(
            final_logits,
            dim=-1,
        )

        .topk(
            top_k,
            dim=-1
        )
    )


    signals = {

        "h_box":
            h_box.unsqueeze(0),

        "delta":
            delta.unsqueeze(0),

        "attn_ent":
            attn_ent.unsqueeze(0),

        "erank":
            erank.unsqueeze(0),

        "topk_probs":
            topk_probs
            .float()
            .unsqueeze(0),

        "topk_ids":
            topk_ids
            .long()
            .unsqueeze(0),
    }


    # IMPORTANT:
    # checkpoint was trained with normalize_h_box=True

    if ckpt_args.get(
        "normalize_h_box",
        True
    ):

        signals["h_box"] = (
            F.normalize(
                signals[
                    "h_box"
                ].float(),
                dim=-1,
            )
        )


    return signals

# ============================================================
# TOKEN-LEVEL HALLUCINATION HTML BUILDER
# ============================================================

def build_highlight_html(token_strings, flagged_set, burst_set, anomaly_scores):
    """
    Build an inline HTML token map.
    Red   = anomaly above threshold (hallucinated)
    Amber = RSNN burst spike
    Hover tooltip shows token index and anomaly score.
    """

    legend = (
        '<div style="font-family:sans-serif;font-size:12px;'
        'margin-bottom:8px;display:flex;gap:14px;flex-wrap:wrap;">'
        '<span>'
        '<span style="background:#ffd8d7;color:#82071e;border-radius:3px;'
        'padding:1px 6px;font-weight:700;">H</span>'
        '&nbsp;Hallucinated (anomaly &gt; threshold)'
        '</span>'
        '<span>'
        '<span style="background:#fff3cd;color:#7d4e00;border-radius:3px;'
        'padding:1px 6px;font-weight:700;">B</span>'
        '&nbsp;RSNN burst spike'
        '</span>'
        '</div>'
    )

    parts = []

    for i, raw in enumerate(token_strings):

        # Humanise the token string
        display = raw
        for src, dst in [
            ('\u0120', ' '),   # Ġ  → space  (GPT-2/Llama BPE)
            ('\u2581', ' '),   # ▁  → space  (SentencePiece)
            ('\u010a', '\n'),  # Ċ  → newline
            ('<0x0A>', '\n'),  # literal newline token
        ]:
            display = display.replace(src, dst)

        # HTML-escape, then convert newlines to <br>
        display = (
            display
            .replace('&', '&amp;')
            .replace('<', '&lt;')
            .replace('>', '&gt;')
            .replace('\n', '<br>')
        )

        if display.strip() == '':
            display = '&nbsp;'

        score = (
            float(anomaly_scores[i])
            if i < len(anomaly_scores)
            else 0.0
        )

        title = f"[{i}] {raw} | anomaly={score:.4f}"

        if i in flagged_set:
            style = (
                "background:#ffd8d7;color:#82071e;"
                "border-radius:3px;padding:1px 3px;cursor:default;"
            )
        elif i in burst_set:
            style = (
                "background:#fff3cd;color:#7d4e00;"
                "border-radius:3px;padding:1px 3px;cursor:default;"
            )
        else:
            style = "cursor:default;"

        parts.append(
            f'<span style="{style}" title="{title}">{display}</span>'
        )

    text_html = (
        '<div style="font-family:sans-serif;font-size:14px;'
        'line-height:1.9;background:#fff;border:1px solid #d0d7de;'
        'border-radius:8px;padding:12px 16px;word-break:break-word;">'
        + ''.join(parts)
        + '</div>'
    )

    heading = (
        '<h4 style="font-family:sans-serif;margin-bottom:6px;'
        'color:#1f2328;">Token-level Hallucination Map</h4>'
    )

    return heading + legend + text_html


# ============================================================
# CHATBOT + RSNN DETECTION
# ============================================================

@spaces.GPU(duration=120)
def respond(
    message,
    history,
    max_tokens,
    temperature,
    top_p,
):

    try:

        # ----------------------------------------------------
        # For now each turn is independent.
        # This matches the training prompt distribution better.
        # ----------------------------------------------------

        prompt = build_prompt(
            message
        )

        prompt_ids = prompt[
            "input_ids"
        ]


        generation_kwargs = {

            "input_ids":
                prompt_ids,

            "attention_mask":
                prompt[
                    "attention_mask"
                ],

            "max_new_tokens":
                int(max_tokens),

            "pad_token_id":
                tok.pad_token_id,

            "do_sample":
                float(
                    temperature
                ) > 0,
        }


        if generation_kwargs[
            "do_sample"
        ]:

            generation_kwargs[
                "temperature"
            ] = float(
                temperature
            )

            generation_kwargs[
                "top_p"
            ] = float(
                top_p
            )


        # ----------------------------------------------------
        # GENERATE ANSWER
        # ----------------------------------------------------

        with torch.no_grad():

            generated = lm.generate(
                **generation_kwargs
            )


        answer_ids = generated[
            :,
            prompt_ids.shape[1]:
        ]


        answer_ids = (
            strip_terminal_tokens(
                answer_ids
            )
        )


        if answer_ids.shape[1] == 0:

            return (
                "The model produced "
                "no scoreable response.",
                "<p><em>No tokens to visualize.</em></p>",
            )


        response_text = tok.decode(

            answer_ids[0],

            skip_special_tokens=True,

            clean_up_tokenization_spaces=False,
        ).strip()


        # ----------------------------------------------------
        # EXTRACT INTERNAL LLM SIGNALS
        # ----------------------------------------------------

        signals = (
            extract_causal_signals(
                prompt_ids,
                answer_ids,
            )
        )


        # ----------------------------------------------------
        # RUN RSNN
        # ----------------------------------------------------

        with torch.no_grad():

            result = rsnn(
                signals
            )


        anomaly = (
            result.anomaly[0]
            .detach()
            .float()
            .cpu()
        )


        rho = (
            result.rho[0]
            .detach()
            .float()
            .cpu()
        )


        burst = (
            result.burst[0]
            .detach()
            .bool()
            .cpu()
        )


        anomaly_threshold = float(
            rsnn.anomaly_star
            .detach()
            .cpu()
        )


        rho_threshold = float(
            rsnn.rho_star
            .detach()
            .cpu()
        )


        # ----------------------------------------------------
        # TOKEN DETECTION
        #
        # Token 0 is excluded because reconstruction anomaly
        # has a structural zero at t=0.
        # ----------------------------------------------------

        anomaly_flags = (
            anomaly
            >
            anomaly_threshold
        )


        burst_flags = (
            burst.clone()
        )


        if anomaly_flags.numel():

            anomaly_flags[0] = False
            burst_flags[0] = False


        flagged_indices = (

            torch.where(
                anomaly_flags
            )[0]

            .tolist()
        )


        burst_indices = (

            torch.where(
                burst_flags
            )[0]

            .tolist()
        )


        token_ids = (

            answer_ids[0]

            .detach()

            .cpu()

            .tolist()
        )


        token_strings = (
            tok.convert_ids_to_tokens(
                token_ids
            )
        )


        # ----------------------------------------------------
        # BUILD TOKEN-LEVEL HTML VISUALIZATION
        # ----------------------------------------------------

        anomaly_list = (
            anomaly
            .detach()
            .cpu()
            .tolist()
        )

        hallu_html = build_highlight_html(
            token_strings,
            flagged_set=set(flagged_indices),
            burst_set=set(burst_indices),
            anomaly_scores=anomaly_list,
        )


        flagged_preview = [

            f"{i}:{token_strings[i]}"

            for i
            in flagged_indices[:12]

            if i
            < len(token_strings)
        ]


        # Ignore structural t=0 score
        valid_anomaly = (

            anomaly[1:]

            if anomaly.numel() > 1

            else anomaly
        )


        valid_rho = (

            rho[1:]

            if rho.numel() > 1

            else rho
        )


        max_anomaly = (

            float(
                valid_anomaly.max()
            )

            if valid_anomaly.numel()

            else float("nan")
        )


        max_rho = (

            float(
                valid_rho.max()
            )

            if valid_rho.numel()

            else float("nan")
        )


        # ----------------------------------------------------
        # SEQUENCE-LEVEL DEMO VERDICT
        # ----------------------------------------------------

        hallucination_detected = (
            len(flagged_indices) > 0
        )


        if hallucination_detected:

            verdict = (
                "⚠️ POTENTIAL "
                "HALLUCINATION DETECTED"
            )

        else:

            verdict = (
                "✅ NO HALLUCINATION "
                "SIGNAL DETECTED"
            )


        preview = (

            ", ".join(
                flagged_preview
            )

            if flagged_preview

            else "none"
        )


        report = f"""

---

### LatticeSpike RSNN Analysis

**Verdict:** {verdict}

**Anomaly threshold:** `{anomaly_threshold:.4f}`  
**Maximum anomaly:** `{max_anomaly:.4f}`  
**Tokens above threshold:** `{len(flagged_indices)}/{len(token_ids)}`  
**RSNN burst tokens:** `{len(burst_indices)}/{len(token_ids)}`  
**Maximum spike density ρ:** `{max_rho:.4f}`  
**ρ threshold:** `{rho_threshold:.4f}`  

**Flagged token positions:** `{preview}`

"""


        return (
            response_text
            +
            "\n\n"
            +
            report,
            hallu_html,
        )


    except Exception as e:

        raise gr.Error(
            "LLM + RSNN inference failed: "
            f"{type(e).__name__}: {e}"
        )


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks() as demo:

    gr.Markdown(
        """
# LatticeSpike Hallucination-Aware Chatbot

**LLM:** Meta Llama 3.2 1B Instruct
**Detector:** LatticeSpike RSNN

The LLM generates an answer, then LatticeSpike analyzes
the internal causal LLM signals token-by-token.

For this first version, each message is evaluated independently.
"""
    )


    hallu_viz = gr.HTML(
        value=(
            "<p style='font-family:sans-serif;color:#57606a;font-size:13px;'>"
            "<em>Submit a question to see the token-level hallucination map.</em>"
            "</p>"
        ),
        label="Token-level Hallucination Visualization",
    )


    chatbot = gr.ChatInterface(

        fn=respond,

        additional_inputs=[

            gr.Slider(
                minimum=16,
                maximum=128,
                value=64,
                step=1,
                label="Max new tokens",
            ),

            gr.Slider(
                minimum=0.0,
                maximum=1.5,
                value=0.7,
                step=0.1,
                label="Temperature",
            ),

            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p",
            ),
        ],

        additional_outputs=[hallu_viz],
    )


demo.queue()


if __name__ == "__main__":

    demo.launch(
        ssr_mode=False
    )
import os
import sys
from pathlib import Path
import gradio as gr
import spaces
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
# ============================================================
# CONFIG
# ============================================================
RSNN_REPO = "saifulsajol/Lattice_Spike"
RSNN_CHECKPOINT = (
    "model_weights/our_model/"
    "latticespike_v5.3_causal_rank1_seed42.pt"
)
LOCAL_RSNN_DIR = "/tmp/Lattice_Spike"
LLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
LLM_REVISION = (
    "9213176726f574b556790deb65791e0c5aa438b6"
)
# Exact system prompt used by the training/evaluation pipeline
SYSTEM_PROMPT = "Answer the question concisely and factually."
PROMPT_MAX_LENGTH = 2048
# ============================================================
# HF TOKEN
# ============================================================
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError(
        "HF_TOKEN is missing. Add HF_TOKEN under "
        "Hugging Face Space Settings -> Secrets."
    )
# ============================================================
# DOWNLOAD LATTICESPIKE REPOSITORY
# ============================================================
print("Downloading LatticeSpike repository...")
snapshot_download(
    repo_id=RSNN_REPO,
    local_dir=LOCAL_RSNN_DIR,
    allow_patterns=[
        "__init__.py",
        "boxes.py",
        "detector.py",
        "llm_signals.py",
        "losses.py",
        "lstm_dyn.py",
        "metrics.py",
        "rsnn.py",
        "signals.py",
        "surrogate.py",
        # Needed because metrics.py imports evaluation.metrics
        "evaluation/*",
        RSNN_CHECKPOINT,
    ],
)
# Needed for:
#   import Lattice_Spike
# and
#   import evaluation
sys.path.insert(
    0,
    str(Path(LOCAL_RSNN_DIR).parent)
)
sys.path.insert(
    0,
    LOCAL_RSNN_DIR
)
from Lattice_Spike import LatticeSpike
# ============================================================
# LOAD RSNN CHECKPOINT
# ============================================================
CHECKPOINT_PATH = os.path.join(