first commit

generate_solutions.py

@@ -0,0 +1,111 @@
+import json
+import os
+import re
+from openai import OpenAI
+
+def load_prompt_data(filepath: str) -> dict:
+    """
+    Loads the prompt data from JSON.
+    """
+    with open(filepath, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+def load_solutions(filepath: str) -> dict:
+    """
+    Loads the existing solutions JSON, or returns a default if file not found.
+    """
+    if os.path.exists(filepath):
+        with open(filepath, "r", encoding="utf-8") as f:
+            return json.load(f)
+    return {}
+
+def save_solutions(filepath: str, solutions: dict):
+    """
+    Saves the solutions dictionary to the solutions.json file (pretty-printed).
+    """
+    with open(filepath, "w", encoding="utf-8") as f:
+        json.dump(solutions, f, indent=4)
+
+def call_LLMs(client, model: str, problem: str, module_header: str) -> str:
+    """
+    Calls the OpenAI chat completion endpoint with the given prompt.
+    """
+    prompt = f"""
+    Here we assume the SystemVerilog is not supported, so don't use the SystemVerilog syntax, such as break statement.
+    Please write a Verilog module that solves the following problem efficiently, using the exact module header below:
+
+    Problem:
+    {problem}
+
+    Module header (must not be changed):
+    {module_header}
+
+    Remember to return only the JSON format:
+    {{
+    "solution": "<verilog code>"
+    }}
+    """
+    
+    try:
+        response = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are a helpful Verilog coding assistant. Please return a JSON object with a key 'solution' containing the Verilog code."},
+                {"role": "user", "content": prompt}
+            ],
+            model=model,
+            max_tokens=3000,
+            temperature=1.5,
+            top_p=0.75,
+        )
+        response_content = response.choices[0].message.content.strip()
+        return response_content
+    except Exception as e:
+        print("Error:", str(e))
+        return json.dumps({"solution": f"Error: {str(e)}"})
+
+def generate_solutions(api_key: str, model_name: str, k: int, prompt_json_file: str = "problems.json", solutions_json_file: str = "solutions.json"):
+    """
+    Generates Verilog solutions for problems using an LLM.
+    """
+    # Initialize OpenAI client
+    client = OpenAI(api_key=api_key)
+    
+    # Load the problem data
+    prompt_data = load_prompt_data(prompt_json_file)
+    
+    # Load or initialize solutions data
+    solutions_data = load_solutions(solutions_json_file)
+
+    if model_name not in solutions_data:
+        solutions_data[model_name] = {}
+
+    for _ in range(k):
+        for category, problems in prompt_data.items():
+            if category not in solutions_data[model_name]:
+                solutions_data[model_name][category] = []
+            
+            for item in problems:
+                problem_statement = item.get("Problem", "")
+                module_header = item.get("Module header", "")
+                module_name = item.get("module")
+                
+                response_json_str = call_LLMs(client, model_name, problem_statement, module_header)
+                response_json_str = response_json_str.strip('`').replace('json', '').replace('```', '')
+                
+                try:
+                    response_json = json.loads(response_json_str)
+                    verilog_code = response_json.get("solution", "")
+                except json.JSONDecodeError:
+                    print(response_json_str)
+                    verilog_code = "Error: Invalid JSON response"
+                
+                print(f"Processing module: {module_name}")
+                category_list = solutions_data[model_name][category]
+                module_entry = next((entry for entry in category_list if entry.get("module") == module_name), None)
+                
+                if module_entry is None:
+                    module_entry = {"module": module_name, "solutions": []}
+                    category_list.append(module_entry)
+                
+                module_entry["solutions"].append({"solution": verilog_code, "pass": ""})
+                save_solutions(solutions_json_file, solutions_data)