#!/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()