import argparse

import json

from pathlib import Path

import torch

import yaml

from peft import PeftModel

from transformers import AutoModelForCausalLM

from encoding_and_decoding.fns import (

encoding_prompt,

decoding_prompt,

decode_code_from_tokens,

)

from utils.tokenization_utils import tokenizer

from utils.config_validation import validate_config

def load_model(config_path: str, lora_adapter_path: str | None = None, device: str = "cuda") -> PeftModel:

cfg_path = Path(config_path)

if not cfg_path.exists():

cfg_path = Path(__file__).parent / config_path

with open(cfg_path) as f:

cfg = yaml.safe_load(f)

# Validate config for model loading

required_keys = [

"model.base_model", "model.dtype",

"logging.save_path",

]

validate_config(cfg, required_keys)

if lora_adapter_path is None:

lora_adapter_path = cfg["logging"]["save_path"]

base_model = cfg["model"]["base_model"]

dtype = getattr(torch, cfg["model"]["dtype"])

# Always load by model id string

model = AutoModelForCausalLM.from_pretrained(base_model, dtype=dtype)

model = PeftModel.from_pretrained(model, lora_adapter_path)

return model.to(device)

def evaluate_encoding(model, cfg) -> float:

"""Legacy helper used during training logs: returns per-bit accuracy only."""

metrics = evaluate_encoding_metrics(model, cfg, num_samples=int(cfg["evaluation"]["encoding_batches"]) * int(cfg["evaluation"]["encoding_batch_size"]))

return metrics["per_bit"]

@torch.no_grad()

def evaluate_encoding_metrics(model, cfg, num_samples: int) -> dict:

"""Compute encoding metrics over a fixed number of samples.

Returns dict with keys: 'per_bit', 'full', 'n'.

"""

bs = int(cfg["evaluation"]["encoding_batch_size"]) or 1

remaining = max(0, int(num_samples))

per_bit_correct = 0

total_bits = 0

full_correct = 0

total = 0

enc_fn = encoding_prompt

while remaining > 0:

cur_bs = min(bs, remaining)

prompts_and_codes = [enc_fn(None, testing=True) for _ in range(cur_bs)]

prompts = [x["tokens"] for x in prompts_and_codes]

codes = [x["code"] for x in prompts_and_codes]

max_len = max(len(x) for x in prompts)

pad_id = tokenizer.pad_token_id

prompts = [[pad_id] * (max_len - len(x)) + x for x in prompts]

input_ids = torch.tensor(prompts, dtype=torch.long, device=model.device)

attention_mask = (input_ids != pad_id).to(model.device)

outputs = model.generate(

input_ids,

do_sample=True,

max_new_tokens=300,

top_p=0.95,

temperature=1.0,

attention_mask=attention_mask,

eos_token_id=tokenizer.eos_token_id,

pad_token_id=pad_id,

)

for j in range(cur_bs):

orig_len = input_ids[j].size(0)

gen = outputs[j, orig_len:]

gen_list = gen.tolist()

if tokenizer.eos_token_id in gen_list:

try:

end = gen_list.index(tokenizer.eos_token_id)

gen_list = gen_list[:end]

except ValueError:

pass

pred_code = decode_code_from_tokens(gen_list, max_bits=int(cfg["data"]["code_bits"]))

target = codes[j]

per_bit_correct += sum(1 for a, b in zip(pred_code, target) if a == b)

total_bits += len(target)

full_correct += 1 if len(pred_code) >= len(target) and all(a == b for a, b in zip(pred_code, target)) else 0

total += 1

remaining -= cur_bs

per_bit = per_bit_correct / total_bits if total_bits else 0.0

full = full_correct / total if total else 0.0

return {"per_bit": per_bit, "full": full, "n": total}

@torch.no_grad()

def evaluate_decoding(model, cfg, messages: list[str]) -> tuple[float, float]:

bs = int(cfg["evaluation"]["decoding_batch_size"])

per_bit_correct = 0

total_bits = 0

full_correct = 0

total = 0

dec_fn = decoding_prompt

for start in range(0, len(messages), bs):

batch_msgs = messages[start : start + bs]

prompts_and_codes = [dec_fn(msg, testing=True) for msg in batch_msgs]

prompts = [d["tokens"] for d in prompts_and_codes]

codes = [d["code"] for d in prompts_and_codes]

max_len = max(len(x) for x in prompts)

pad_id = tokenizer.pad_token_id

prompts = [[pad_id] * (max_len - len(x)) + x for x in prompts]

input_ids = torch.tensor(prompts, dtype=torch.long, device=model.device)

attention_mask = (input_ids != pad_id).to(model.device)

outputs = model.generate(

input_ids,

do_sample=False,

max_new_tokens=20,

top_p=None,

temperature=None,

attention_mask=attention_mask,

eos_token_id=tokenizer.eos_token_id,

pad_token_id=pad_id,

)

for orig_code, out_ids, inp in zip(codes, outputs, input_ids):

inp_len = inp.size(0)

# take only tokens up to eos if present

gen_ids = out_ids[inp_len:]

if tokenizer.eos_token_id in gen_ids:

eos_pos = (gen_ids == tokenizer.eos_token_id).nonzero(as_tuple=True)[0][0].item()

gen_ids = gen_ids[:eos_pos]

gen_list = gen_ids.tolist()

gen_str = tokenizer.decode(gen_list)

# extract only the first N digits (0/1) from the generated text

code_bits = int(cfg["data"]["code_bits"]) if "data" in cfg and "code_bits" in cfg["data"] else 16

pred_digits = []

for ch in gen_str:

if ch in ("0", "1"):

pred_digits.append(int(ch))

if len(pred_digits) >= code_bits:

break

correct_bits = sum(1 for o, d in zip(orig_code, pred_digits) if o == d)

per_bit_correct += correct_bits

total_bits += len(orig_code)

full_correct += 1 if len(pred_digits) >= len(orig_code) and correct_bits == len(orig_code) else 0

total += 1

return (full_correct / total if total else 0.0, per_bit_correct / total_bits if total_bits else 0.0)

@torch.no_grad()

def evaluate_decoding_on_self_generated(model, cfg, num_samples: int) -> tuple[float, float]:

"""Evaluate decoding on messages first encoded by the model itself.

Pipeline per mini-batch:

1) Generate encoded messages from encoding prompts (fresh random codes).

2) Build decoding prompts using those generated messages (pretokenized).

3) Run decoding and compare to the original intended codes.

"""

bs = int(cfg["evaluation"]["decoding_batch_size"]) or 1

remain = max(0, int(num_samples))

per_bit_correct = 0

total_bits = 0

full_correct = 0

total = 0

pad_id = tokenizer.pad_token_id

code_bits = int(cfg["data"]["code_bits"]) if "data" in cfg and "code_bits" in cfg["data"] else 16

enc_fn = encoding_prompt

dec_fn = decoding_prompt

while remain > 0:

cur_bs = min(bs, remain)

# Step 1: create encoding prompts and generate encoded messages

enc_batch = [enc_fn(None, testing=True) for _ in range(cur_bs)]

enc_prompts = [e["tokens"] for e in enc_batch]

enc_codes = [e["code"] for e in enc_batch]

max_len_enc = max(len(x) for x in enc_prompts)

enc_prompts = [[pad_id] * (max_len_enc - len(x)) + x for x in enc_prompts]

enc_input_ids = torch.tensor(enc_prompts, dtype=torch.long, device=model.device)

enc_attention_mask = (enc_input_ids != pad_id).to(model.device)

enc_outputs = model.generate(

enc_input_ids,

do_sample=True,

max_new_tokens=300,

top_p=0.95,

temperature=1.0,

attention_mask=enc_attention_mask,

eos_token_id=tokenizer.eos_token_id,

pad_token_id=pad_id,

)

generated_msgs: list[list[int]] = []

for j in range(cur_bs):

orig_len = enc_input_ids[j].size(0)

gen = enc_outputs[j, orig_len:]

gen_list = gen.tolist()

if tokenizer.eos_token_id in gen_list:

try:

end = gen_list.index(tokenizer.eos_token_id)

gen_list = gen_list[:end]

except ValueError:

pass

generated_msgs.append(gen_list)

# Step 2: build decoding prompts using generated messages (pretokenized)

dec_prompts_and_codes = [

dec_fn(msg_tokens, code=enc_codes[i], testing=True, pretokenized_message=True)

for i, msg_tokens in enumerate(generated_msgs)

]

dec_prompts = [d["tokens"] for d in dec_prompts_and_codes]

max_len_dec = max(len(x) for x in dec_prompts)

dec_prompts = [[pad_id] * (max_len_dec - len(x)) + x for x in dec_prompts]

dec_input_ids = torch.tensor(dec_prompts, dtype=torch.long, device=model.device)

dec_attention_mask = (dec_input_ids != pad_id).to(model.device)

# Step 3: generate decoding outputs deterministically and score

dec_outputs = model.generate(

dec_input_ids,

do_sample=False,

max_new_tokens=20,

top_p=None,

temperature=None,

attention_mask=dec_attention_mask,

eos_token_id=tokenizer.eos_token_id,

pad_token_id=pad_id,

)

for gt_code, out_ids, inp in zip(enc_codes, dec_outputs, dec_input_ids):

inp_len = inp.size(0)

gen_ids = out_ids[inp_len:]

if tokenizer.eos_token_id in gen_ids:

eos_pos = (gen_ids == tokenizer.eos_token_id).nonzero(as_tuple=True)[0][0].item()

gen_ids = gen_ids[:eos_pos]

gen_str = tokenizer.decode(gen_ids)

pred_digits: list[int] = []

for ch in gen_str:

if ch in ("0", "1"):

pred_digits.append(int(ch))

if len(pred_digits) >= code_bits:

break

correct_bits = sum(1 for o, d in zip(gt_code, pred_digits) if o == d)

per_bit_correct += correct_bits

total_bits += len(gt_code)

full_correct += 1 if len(pred_digits) >= len(gt_code) and correct_bits == len(gt_code) else 0

total += 1

remain -= cur_bs

return (

full_correct / total if total else 0.0,

per_bit_correct / total_bits if total_bits else 0.0,

)

@torch.no_grad()

def evaluate_o3_messages(model, cfg, codes_per_message: int = 80) -> dict:

with open(cfg["data"]["o3_messages_path"], "r") as f:

messages_raw = json.load(f)

messages = [m["message"] for m in messages_raw]

# Replicate each message 'codes_per_message' times so each gets a fresh random code

if codes_per_message > 1:

messages = [m for m in messages for _ in range(int(codes_per_message))]

full, per_bit = evaluate_decoding(model, cfg, messages)

return {"full": full, "per_bit": per_bit, "n": len(messages)}

@torch.no_grad()

def sample_encoding_examples(model, cfg, num: int = 6, code: list[int] | None = None):

# Allow passing a single code or a list of codes (one per example)

enc_fn = encoding_prompt

if code is None:

prompts_and_codes = [enc_fn(None, testing=True) for _ in range(num)]

else:

# If 'code' looks like a flat list of ints, use it for all examples

if len(code) == 0 or isinstance(code[0], int): # type: ignore[index]

prompts_and_codes = [enc_fn(None, testing=True, code=code) for _ in range(num)] # type: ignore[arg-type]

else:

# Assume list[list[int]]; length must match 'num'

codes_list = code # type: ignore[assignment]

if len(codes_list) != num:

raise ValueError(f"When providing multiple codes, expected {num} codes, got {len(codes_list)}")

prompts_and_codes = [enc_fn(None, testing=True, code=codes_list[i]) for i in range(num)] # type: ignore[index]

prompts = [x["tokens"] for x in prompts_and_codes]

codes = [x["code"] for x in prompts_and_codes]

max_len = max(len(x) for x in prompts)

pad_id = tokenizer.pad_token_id

prompts = [[pad_id] * (max_len - len(x)) + x for x in prompts]

input_ids = torch.tensor(prompts, dtype=torch.long, device=model.device)

attention_mask = (input_ids != pad_id).to(model.device)

outputs = model.generate(

input_ids,

do_sample=True,

max_new_tokens=200,

top_p=0.95,

temperature=1.0,

attention_mask=attention_mask,

eos_token_id=tokenizer.eos_token_id,

pad_token_id=pad_id,

)

examples = []

for j in range(num):

orig_len = input_ids[j].size(0)

gen = outputs[j, orig_len:]

gen_list = gen.tolist()

if tokenizer.eos_token_id in gen_list:

try:

end = gen_list.index(tokenizer.eos_token_id)

gen_list = gen_list[:end]

except ValueError:

pass

pred_bits = decode_code_from_tokens(gen_list, max_bits=int(cfg["data"]["code_bits"]))

pred_code = "".join(str(b) for b in pred_bits)

target_code = "".join(str(b) for b in codes[j])

gen_text = tokenizer.decode(gen_list)

# Build token-by-token wrapped representation, using per-token decode

token_pieces = [tokenizer.decode([tid]) for tid in gen_list]

wrapped = "".join(f"[{piece}]" for piece in token_pieces)

examples.append({

"pred_code": pred_code,

"target_code": target_code,

"match": pred_code == target_code,

"generated_text": gen_text,

"generated_tokens_wrapped": wrapped,

})

return examples

@torch.no_grad()

def sample_decoding_examples(model, cfg, messages: list[str]):

dec_fn = decoding_prompt

prompts_and_codes = [dec_fn(msg, testing=True) for msg in messages]

prompts = [d["tokens"] for d in prompts_and_codes]

codes = [d["code"] for d in prompts_and_codes]

max_len = max(len(x) for x in prompts)

pad_id = tokenizer.pad_token_id

prompts = [[pad_id] * (max_len - len(x)) + x for x in prompts]

input_ids = torch.tensor(prompts, dtype=torch.long, device=model.device)

attention_mask = (input_ids != pad_id).to(model.device)

outputs = model.generate(

input_ids,

do_sample=False,

max_new_tokens=20,

top_p=None,

temperature=None,

attention_mask=attention_mask,

eos_token_id=tokenizer.eos_token_id,

pad_token_id=pad_id,

)

examples = []

for msg, code, out_ids, inp in zip(messages, codes, outputs, input_ids):

inp_len = inp.size(0)

gen_ids = out_ids[inp_len:]

if tokenizer.eos_token_id in gen_ids:

eos_pos = (gen_ids == tokenizer.eos_token_id).nonzero(as_tuple=True)[0][0].item()

gen_ids = gen_ids[:eos_pos]

gen_str = tokenizer.decode(gen_ids)

# extract only 0/1 digits up to code_bits

code_bits = int(cfg["data"]["code_bits"]) if "data" in cfg and "code_bits" in cfg["data"] else 16

pred_digits = []

for ch in gen_str:

if ch in ("0", "1"):

pred_digits.append(ch)

if len(pred_digits) >= code_bits:

break

pred_code = "".join(pred_digits)

target_code = "".join(str(x) for x in code)

examples.append({

"message": msg,

"pred_code": pred_code,

"target_code": target_code,

"match": pred_code == target_code,

})

return examples

def main():

parser = argparse.ArgumentParser(description="Evaluate enc_and_dec model")

parser.add_argument("--config", type=str, default="config.yaml")

parser.add_argument("--o3-only", action="store_true")

args = parser.parse_args()

cfg_path = args.config

model = load_model(cfg_path)

with open(cfg_path) as f:

cfg = yaml.safe_load(f)

# Validate config for evaluation

required_keys = [

"data.llama_messages_path", "data.o3_messages_path",

"data.min_len", "data.max_len", "data.code_bits",

"evaluation.final.o3_messages",

"evaluation.encoding_batch_size", "evaluation.decoding_batch_size",

]

validate_config(cfg, required_keys)

if args.o3_only:

o3 = evaluate_o3_messages(model, cfg)

print(f"O3 messages: full={o3['full']*100:.1f}% | per-bit={o3['per_bit']*100:.1f}%")

return

enc_acc = evaluate_encoding(model, cfg)

# small sample of llama messages

with open(cfg["data"]["llama_messages_path"], "r") as f:

messages_raw = json.load(f)

messages = [m["message"] for m in messages_raw if cfg["data"]["min_len"] < m["length"] < cfg["data"]["max_len"]]

dec_full, dec_pbit = evaluate_decoding(model, cfg, messages[:300])

print(f"Encoding per-bit: {enc_acc*100:.1f}% | Decoding full: {dec_full*100:.1f}% | per-bit: {dec_pbit*100:.1f}%")

if cfg["evaluation"]["final"]["o3_messages"]:

o3 = evaluate_o3_messages(model, cfg)

print(f"O3 messages: full={o3['full']*100:.1f}% | per-bit={o3['per_bit']*100:.1f}%")

if __name__ == "__main__":

main()