import json
from pathlib import Path

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

REPO_ID = "scale-lab/CovR-Qwen2.5-Coder-7B-Instruct"
BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"

# Authenticate first with `hf auth login` while this repository is private.
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto"
)
model = PeftModel.from_pretrained(base, REPO_ID)
model.eval()

prompts = json.loads(Path(hf_hub_download(REPO_ID, "prompt_templates.json")).read_text())
spec = "A 4-bit counter. On each rising clock edge, synchronous active-high reset clears count; otherwise enable increments count modulo 16."
rtl = """module counter(input logic clk, rst, en, output logic [3:0] count);
  always_ff @(posedge clk) begin
    if (rst) count <= 4'd0;
    else if (en) count <= count + 4'd1;
  end
endmodule"""
messages = [
    {"role": "system", "content": prompts["system"]},
    {"role": "user", "content": prompts["user_template"].format(
        natural_spec=spec, rtl_code=rtl, context=""
    )},
]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, tokenize=True,
    return_dict=True, return_tensors="pt",
).to(model.device)
with torch.inference_mode():
    outputs = model.generate(
        **inputs, max_new_tokens=4096, do_sample=True,
        temperature=0.1, top_p=1.0, top_k=50,
        pad_token_id=tokenizer.eos_token_id,
    )
print(tokenizer.decode(outputs[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))