| 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" |
|
|
| |
| 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)) |
|
|
