add readme

README.md

@@ -1,3 +1,141 @@
 ---
 license: gpl-3.0
+base_model: deepseek-ai/deepseek-coder-7b-instruct-v1.5
+library_name: peft 
+
 ---
+
+# OriGen: Enhancing RTL Code Generation with Code-to-Code Augmentation and Self-Reflection
+
+### Introduction
+OriGen is a fine-tuned lora model designed for Verilog code generation. It is trained on top of DeepSeek Coder 7B using datasets generated from code-to-code augmentation and self-reflection.
+
+OriGen_Fix is a fine-tuned lora model designed for fixing syntax errors in Verilog code. It is trained based on OriGen. 
+
+
+The models have been uploaded to Hugging Face, and the repository contains the inference scripts. The dataset and data generation flow will be released soon.
+
+- **Huggingface**: 
+    - https://huggingface.co/henryen/OriGen
+    - https://huggingface.co/henryen/OriGen_Fix
+- **Repository**: https://github.com/pku-liang/OriGen
+
+### Evaluation Results
+<img src="figures/evaluation.png" alt="evaluation" width="1000"/>
+
+### Quick Start
+
+Before running the following code, please install the required packages:
+
+```bash
+conda create -n origen python=3.11
+conda activate origen
+pip install -r requirements.txt
+```
+
+Here is an example of how to use the model. Please note that the base model, DeepSeek Coder 7B, is loaded in float16 precision, even though its default precision is bfloat16. This choice was made because our experiments showed that Lora trained in float16 outperforms those trained in bfloat16.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+import torch
+from peft import PeftModel
+import json
+
+model_name = "deepseek-ai/deepseek-coder-7b-instruct-v1.5"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    low_cpu_mem_usage=True,
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2",
+    device_map="auto",
+).to("cuda")
+
+model = PeftModel.from_pretrained(model, model_id="henryen/OriGen_Fix")
+model.eval()
+
+streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+prompt_template = """
+        ### Instruction: Please act as a professional Verilog designer. Your task is to debug a Verilog module.\nYou will receive a task, an original verilog code with syntax and function errors, and the corresponding error messages. \nYou should generate a corrected code based on the original code and the error messages
+        Your task:
+        {description}
+        Original code:
+        {original_code}
+        Error message:
+        {error}
+        You should now generate a correct code.
+        ### Response:{header}
+        """
+
+def generate_code(data):
+    description = data["description"]
+    original_code = data["original_code"]
+    error = data["error"]
+    header = data["module_header"]
+    prompt = prompt_template.format(description=description, original_code=original_code, error=error, header=header)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+    outputs = model.generate(
+        **inputs, 
+        max_new_tokens=1000, 
+        do_sample=False, 
+        temperature=0, 
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        streamer=streamer
+    )
+    
+    input_length = len(inputs[0])
+    completion = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+    return completion
+
+input_file = "./data/example/example_error.jsonl"
+output_file = "./data/example/example_fix.jsonl"
+
+with open(input_file, "r") as f, open(output_file, "w") as f2:
+    for line in f:
+        data = json.loads(line)
+        completion = generate_code(data)
+        json.dump({"task_id": data["task_id"], "completion": completion}, f2)
+        f2.write("\n")
+```
+
+
+The output will be:
+```verilog
+    wire and0_out;
+    wire and1_out;
+    wire or0_out;
+
+    assign and0_out = a & b;
+    assign and1_out = c & d;
+    assign or0_out = and0_out | and1_out;
+    assign out = or0_out;
+    assign out_n = ~or0_out;
+
+endmodule
+```
+
+You can check its correctness using testbench provided under the folder `./data/example/`.
+```bash
+cd ./data/example/
+evaluate_functional_correctness example_fix.jsonl --problem_file example_problem.jsonl
+```
+
+
+### Paper
+**Arxiv:** https://arxiv.org/abs/2407.16237
+
+Please cite our paper if you find this model useful.
+
+```
+@article{2024origen,
+  title={OriGen: Enhancing RTL Code Generation with Code-to-Code Augmentation and Self-Reflection},
+  author={Cui, Fan and Yin, Chenyang and Zhou, Kexing and Xiao, Youwei and Sun, Guangyu and Xu, Qiang and Guo, Qipeng and Song, Demin and Lin, Dahua and Zhang, Xingcheng and others},
+  journal={arXiv preprint arXiv:2407.16237},
+  year={2024}
+}
+```

data/example/example_error.jsonl

@@ -0,0 +1 @@
+{"task_id": "wire_decl", "description": "Implement the following circuit. Create two intermediate wires (named anything you want) to connect the AND and OR gates together. Note that the wire that feeds the NOT gate is really wire `out`, so you do not necessarily need to declare a third wire here. Notice how wires are driven by exactly one source (output of a gate), but can feed multiple inputs.\n\n// The circuit is composed of two layers. The first layer, counting from the input, is two AND gates: one whose input is connected to a and b, and the second is connected to c and d. The second layer there is an OR gate to OR the two AND outputs, connected the output 'out'. Additionally, there is an inverted output 'out_n'.", "module_header": "module top_module (\n\tinput a,\n\tinput b,\n\tinput c,\n\tinput d,\n\toutput out,\n\toutput out_n );\n", "original_code": "module top_module (\n\tinput a,\n\tinput b,\n\tinput c,\n\tinput d,\n\toutput out,\n\toutput out_n );\n wire and0_out;\n         wire and1_out;\n         wire or0_out;\n\n         and and0 (\n           .a(a),\n           .b(b),\n           .o(and0_out)\n         );\n         and and1 (\n           .a(c),\n           .b(d),\n           .o(and1_out)\n         );\n         or or0 (\n           .a(and0_out),\n           .b(and1_out),\n           .o(or0_out)\n         );\n         assign out = or0_out;\n         assign out_n = ~or0_out;\n       endmodule", "error": "17: error: Gates do not have port names.\n22: error: Gates do not have port names.\n12: error: Gates do not have port names.\n"}

data/example/example_fix.jsonl

@@ -0,0 +1 @@
+{"task_id": "wire_decl", "completion": " wire and0_out;\n         wire and1_out;\n         wire or0_out;\n\n         assign and0_out = a & b;\n         assign and1_out = c & d;\n         assign or0_out = and0_out | and1_out;\n         assign out = or0_out;\n         assign out_n = ~or0_out;\n\nendmodule\n"}

data/example/example_fix.jsonl_results.jsonl

@@ -0,0 +1 @@
+{"task_id": "wire_decl", "completion": " wire and0_out;\n         wire and1_out;\n         wire or0_out;\n\n         assign and0_out = a & b;\n         assign and1_out = c & d;\n         assign or0_out = and0_out | and1_out;\n         assign out = or0_out;\n         assign out_n = ~or0_out;\n\nendmodule\n", "result": "passed", "passed": true}

data/example/example_problem.jsonl

@@ -0,0 +1 @@
+{"task_id": "wire_decl", "prompt": "module top_module (\n\tinput a,\n\tinput b,\n\tinput c,\n\tinput d,\n\toutput out,\n\toutput out_n );\n", "canonical_solution": "\t\n\twire w1, w2;\n\tassign w1 = a&b;\n\tassign w2 = c&d;\n\tassign out = w1|w2;\n\tassign out_n = ~out;\n\t\nendmodule\n", "test": "`timescale 1 ps/1 ps\n`define OK 12\n`define INCORRECT 13\n// hdlbits_prop {len: 5}\nmodule reference_module (\n\tinput a,\n\tinput b,\n\tinput c,\n\tinput d,\n\toutput out,\n\toutput out_n );\n\t\n\twire w1, w2;\n\tassign w1 = a&b;\n\tassign w2 = c&d;\n\tassign out = w1|w2;\n\tassign out_n = ~out;\n\t\nendmodule\n\n\nmodule stimulus_gen (\n\tinput clk,\n\toutput reg a,b,c,d,\n\toutput reg[511:0] wavedrom_title,\n\toutput reg wavedrom_enable\n);\n\n\n// Add two ports to module stimulus_gen:\n//    output [511:0] wavedrom_title\n//    output reg wavedrom_enable\n\n\ttask wavedrom_start(input[511:0] title = \"\");\n\tendtask\n\t\n\ttask wavedrom_stop;\n\t\t#1;\n\tendtask\t\n\n\n\n\tinitial begin\n\t\t{a,b,c,d} = 4'h0;\n\t\t@(negedge clk);\n\t\twavedrom_start(\"Exhaustive test\");\n\t\trepeat(20) @(posedge clk, negedge clk)\n\t\t\t{d,c,b,a} <= {d,c,b,a} + 1'b1;\n\t\twavedrom_stop();\n\t\trepeat(100) @(posedge clk, negedge clk) begin\n\t\t\t{a,b,c,d} <= $random;\n\t\tend\n\t\t\n\t\t#1 $finish;\n\tend\n\t\nendmodule\n\nmodule tb();\n\n\ttypedef struct packed {\n\t\tint errors;\n\t\tint errortime;\n\t\tint errors_out;\n\t\tint errortime_out;\n\t\tint errors_out_n;\n\t\tint errortime_out_n;\n\n\t\tint clocks;\n\t} stats;\n\t\n\tstats stats1;\n\t\n\t\n\twire[511:0] wavedrom_title;\n\twire wavedrom_enable;\n\tint wavedrom_hide_after_time;\n\t\n\treg clk=0;\n\tinitial forever\n\t\t#5 clk = ~clk;\n\n\tlogic a;\n\tlogic b;\n\tlogic c;\n\tlogic d;\n\tlogic out_ref;\n\tlogic out_dut;\n\tlogic out_n_ref;\n\tlogic out_n_dut;\n\n\tinitial begin \n\t\t$dumpfile(\"wave.vcd\");\n\t\t$dumpvars(1, stim1.clk, tb_mismatch ,a,b,c,d,out_ref,out_dut,out_n_ref,out_n_dut );\n\tend\n\n\n\twire tb_match;\t\t// Verification\n\twire tb_mismatch = ~tb_match;\n\t\n\tstimulus_gen stim1 (\n\t\t.clk,\n\t\t.* ,\n\t\t.a,\n\t\t.b,\n\t\t.c,\n\t\t.d );\n\treference_module good1 (\n\t\t.a,\n\t\t.b,\n\t\t.c,\n\t\t.d,\n\t\t.out(out_ref),\n\t\t.out_n(out_n_ref) );\n\t\t\n\ttop_module top_module1 (\n\t\t.a,\n\t\t.b,\n\t\t.c,\n\t\t.d,\n\t\t.out(out_dut),\n\t\t.out_n(out_n_dut) );\n\n\t\n\tbit strobe = 0;\n\ttask wait_for_end_of_timestep;\n\t\trepeat(5) begin\n\t\t\tstrobe <= !strobe;  // Try to delay until the very end of the time step.\n\t\t\t@(strobe);\n\t\tend\n\tendtask\t\n\n\t\n\tfinal begin\n\t\tif (stats1.errors_out) $display(\"Hint: Output '%s' has %0d mismatches. First mismatch occurred at time %0d.\", \"out\", stats1.errors_out, stats1.errortime_out);\n\t\telse $display(\"Hint: Output '%s' has no mismatches.\", \"out\");\n\t\tif (stats1.errors_out_n) $display(\"Hint: Output '%s' has %0d mismatches. First mismatch occurred at time %0d.\", \"out_n\", stats1.errors_out_n, stats1.errortime_out_n);\n\t\telse $display(\"Hint: Output '%s' has no mismatches.\", \"out_n\");\n\n\t\t$display(\"Hint: Total mismatched samples is %1d out of %1d samples\\n\", stats1.errors, stats1.clocks);\n\t\t$display(\"Simulation finished at %0d ps\", $time);\n\t\t$display(\"Mismatches: %1d in %1d samples\", stats1.errors, stats1.clocks);\n\tend\n\t\n\t// Verification: XORs on the right makes any X in good_vector match anything, but X in dut_vector will only match X.\n\tassign tb_match = ( { out_ref, out_n_ref } === ( { out_ref, out_n_ref } ^ { out_dut, out_n_dut } ^ { out_ref, out_n_ref } ) );\n\t// Use explicit sensitivity list here. @(*) causes NetProc::nex_input() to be called when trying to compute\n\t// the sensitivity list of the @(strobe) process, which isn't implemented.\n\talways @(posedge clk, negedge clk) begin\n\n\t\tstats1.clocks++;\n\t\tif (!tb_match) begin\n\t\t\tif (stats1.errors == 0) stats1.errortime = $time;\n\t\t\tstats1.errors++;\n\t\tend\n\t\tif (out_ref !== ( out_ref ^ out_dut ^ out_ref ))\n\t\tbegin if (stats1.errors_out == 0) stats1.errortime_out = $time;\n\t\t\tstats1.errors_out = stats1.errors_out+1'b1; end\n\t\tif (out_n_ref !== ( out_n_ref ^ out_n_dut ^ out_n_ref ))\n\t\tbegin if (stats1.errors_out_n == 0) stats1.errortime_out_n = $time;\n\t\t\tstats1.errors_out_n = stats1.errors_out_n+1'b1; end\n\n\tend\nendmodule\n"}