File size: 1,669 Bytes
d8b60e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | 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))
|