inference_aes.py

#!/usr/bin/env python3
"""
AES Security IP Inference Script for Elinnos AES LoRA Adapter.

Loads the all-merged base model (Qwen2.5-7B + V1+V2+V3+V4+SRAM+I2CS baked in)
and applies the AES LoRA adapter for inference.

Configuration:
  - Temperature: 0.2 (low randomness, deterministic-ish RTL generation)
  - Max new tokens: 8192 (enough for full AES RTL files)
  - System prompt: AES security IP persona from training dataset

Usage:
    python3 inference_aes.py --interactive
    python3 inference_aes.py --prompt "Give me the aes_top.sv RTL"
    python3 inference_aes.py --adapter-path /path/to/aes_lora
    python3 inference_aes.py --base-path /path/to/all_merged
"""
import argparse
import re
import sys
import time
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

WORKSPACE = Path("/workspace/elinnos")
DEFAULT_BASE = WORKSPACE / "merged_models" / "elinnos_all_merged_final"
DEFAULT_ADAPTER = WORKSPACE / "elinnos-qwen2.5-7b-aes-lora"
CHAT_TEMPLATE_SRC = WORKSPACE / "elinnos-qwen2.5-7b-multi-ip-lora-v4" / "chat_template.jinja"

TEMPERATURE = 0.2
MAX_NEW_TOKENS = 8192

SYSTEM_PROMPT = """You are Elinnos, a hardware design assistant specialising in SystemVerilog / Verilog RTL
and verification for Elinnos IP blocks.

When generating AES / security IP artifacts:
- Follow the modular AES structure (aes_top with configurable APB or AHB-Lite slave
  interface, aes_core, aes_regfile, aes_ctrl, aes_key_expand, aes_round_core,
  SubBytes/ShiftRows/MixColumns and inverse transforms, S-Box tables).
- Support AES-128 and AES-256 encrypt/decrypt with iterative (default) or optional
  pipelined round architecture via compile-time macros (AES_IF_APB, AES_IF_AHB,
  AES_ITERATIVE, AES_PIPELINED).
- Host interface is selected at compile time (APB default or AHB-Lite).
- Use `timescale 1ns/1ps in testbench files.
- Return only the requested file content unless asked for a manifest.
- Preserve naming: aes_* modules, register map (CTRL/STATUS/KEY*/DATA_*/IRQ_*),
  and directed TB stimulus/check tasks."""

_DOLLAR_TAG_RE = re.compile(r'([A-Za-z_][A-Za-z0-9_]*)\$\$([A-Za-z0-9]+)')


def normalize_dollar_tags(text: str) -> str:
    first_tag: dict[str, str] = {}

    def _repl(m: "re.Match[str]") -> str:
        base, tag = m.group(1), m.group(2)
        canonical = first_tag.setdefault(base, tag)
        return f"{base}$${canonical}"

    return _DOLLAR_TAG_RE.sub(_repl, text)


def strip_dollar_tags(text: str) -> str:
    return _DOLLAR_TAG_RE.sub(r"\1", text)


def parse_args():
    p = argparse.ArgumentParser(description="Elinnos AES Security IP LoRA inference")
    p.add_argument("--base-path", type=str, default=str(DEFAULT_BASE))
    p.add_argument("--adapter-path", type=str, default=str(DEFAULT_ADAPTER))
    p.add_argument("--prompt", type=str, default=None)
    p.add_argument("--interactive", action="store_true")
    p.add_argument("--temperature", type=float, default=TEMPERATURE)
    p.add_argument("--max-new-tokens", type=int, default=MAX_NEW_TOKENS)
    p.add_argument("--system-prompt", type=str, default=None)
    p.add_argument("--save-output", type=str, default=None)
    p.add_argument("--strip-tags", action="store_true")
    return p.parse_args()


def load_model(base_path, adapter_path):
    print(f"Loading tokenizer from {base_path}...")
    tokenizer = AutoTokenizer.from_pretrained(str(base_path), trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    chat_template_path = Path(adapter_path) / "chat_template.jinja"
    if not chat_template_path.exists():
        chat_template_path = Path(base_path) / "chat_template.jinja"
    if not chat_template_path.exists():
        chat_template_path = CHAT_TEMPLATE_SRC
    if chat_template_path.exists():
        tokenizer.chat_template = chat_template_path.read_text()
        print(f"  Chat template loaded from {chat_template_path}")

    print(f"Loading model from {base_path} (bf16)...")
    model = AutoModelForCausalLM.from_pretrained(
        str(base_path),
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )

    if adapter_path and Path(adapter_path).is_dir():
        print(f"Applying LoRA adapter from {adapter_path}...")
        model = PeftModel.from_pretrained(model, str(adapter_path))
    model.eval()

    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_mem = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        print(f"  GPU: {gpu_name} ({gpu_mem:.1f} GB)")

    print("Model ready.\n")
    return model, tokenizer


def generate_response(model, tokenizer, messages, temperature, max_new_tokens):
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    start = time.time()
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=temperature > 0,
            top_p=0.9,
            repetition_penalty=1.05,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    elapsed = time.time() - start

    input_len = inputs["input_ids"].shape[1]
    generated = outputs[0][input_len:]
    response = tokenizer.decode(generated, skip_special_tokens=True)
    response = normalize_dollar_tags(response)
    n_tokens = len(generated)
    tps = n_tokens / elapsed if elapsed > 0 else 0
    return response, n_tokens, elapsed, tps


def main():
    args = parse_args()
    system_prompt = args.system_prompt if args.system_prompt else SYSTEM_PROMPT

    print("=" * 70)
    print("  ELINNOS AES SECURITY IP INFERENCE")
    print(f"  Base: {args.base_path}")
    print(f"  Adapter: {args.adapter_path}")
    print(f"  Temperature: {args.temperature}")
    print(f"  Max new tokens: {args.max_new_tokens}")
    print("=" * 70 + "\n")

    if args.prompt:
        model, tokenizer = load_model(args.base_path, args.adapter_path)
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": args.prompt},
        ]
        print(f"User: {args.prompt}\n")
        response, n_tokens, elapsed, tps = generate_response(
            model, tokenizer, messages,
            temperature=args.temperature,
            max_new_tokens=args.max_new_tokens,
        )
        clean_response = strip_dollar_tags(response)
        print(f"Assistant ({n_tokens} tokens, {elapsed:.1f}s, {tps:.1f} tok/s):\n")
        print(clean_response)
        print(f"\n{'─' * 60}")

        if args.save_output:
            Path(args.save_output).write_text(clean_response)
            print(f"Output saved to: {args.save_output}")
        return

    if not args.interactive:
        print("No prompt provided. Use --prompt or --interactive.")
        print("Example:")
        print('  python3 inference_aes.py --interactive')
        print('  python3 inference_aes.py --prompt "Give me the aes_top.sv RTL"')
        return

    model, tokenizer = load_model(args.base_path, args.adapter_path)

    print("=" * 70)
    print("  INTERACTIVE MODE -- type 'exit' or 'quit' to stop")
    print("=" * 70 + "\n")

    conversation = [{"role": "system", "content": system_prompt}]

    while True:
        try:
            user_input = input("User> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break

        if user_input.lower() in ("exit", "quit"):
            print("Exiting.")
            break

        if not user_input:
            continue

        conversation.append({"role": "user", "content": user_input})
        response, n_tokens, elapsed, tps = generate_response(
            model, tokenizer, conversation,
            temperature=args.temperature,
            max_new_tokens=args.max_new_tokens,
        )
        clean_response = strip_dollar_tags(response)
        print(f"\nAssistant ({n_tokens} tokens, {elapsed:.1f}s, {tps:.1f} tok/s):\n")
        print(clean_response)
        print()
        conversation.append({"role": "assistant", "content": response})

        if args.save_output:
            with open(args.save_output, "a") as f:
                f.write(f"User: {user_input}\n\nAssistant: {clean_response}\n\n{'='*70}\n\n")


if __name__ == "__main__":
    main()