update template.py

template.py

@@ -3,23 +3,37 @@ from datasets import load_dataset
 from openai import OpenAI
 from together import Together
 
-# --- Dataset Configuration (NOT TO BE CHANGED) ---
+# === DATASET CONFIGURATION (DO NOT MODIFY) ===
 GT_DATASET_NAME = "kostis-init/CP-Bench"
 DATASET_SPLIT = "train"
 PROBLEM_ID_COLUMN = "id"
 PROBLEM_DESCRIPTION_COLUMN = "description"
 PROBLEM_DATA_COLUMN = "input_data"
 PROBLEM_DECISION_VARS_COLUMN = "decision_variables"
-# --- End of Dataset Configuration ---
+# ==============================================
 
 #######################################################################
-# This is an example script to generate constraint models using LLMs. #
+# Template script to generate constraint models using LLMs.           #
 # You can use this as a starting point for your own approach.         #
 #######################################################################
 
-# --- Submission Approach Configuration ---
-LLM_CLIENT = OpenAI(api_key='YOUR_API_KEY')  # TODO: Replace with your OpenAI API key, and/or use another LLM client (e.g. Together AI)
-LLM_ID = "gpt-4o"  # TODO: Choose the LLM model you want to use
+# === CHOOSE LLM CLIENT AND MODEL CONFIGURATION ===
+
+# Example 1: OpenAI (e.g., GPT-4o)
+# LLM_CLIENT = OpenAI(api_key="YOUR_API_KEY")
+# LLM_ID = "gpt-4o"
+
+# Example 2: DeepSeek (or any other OpenAI-compatible API)
+# LLM_CLIENT = OpenAI(api_key="DEEPSEEK_API_KEY", base_url="https://api.deepseek.com")
+# LLM_ID = "deepseek-chat"
+
+# Example 3: Together.ai
+# LLM_CLIENT = Together(api_key="TOGETHER_API_KEY")
+# LLM_ID = "mistralai/Mixtral-8x22B-Instruct-v0.1"
+
+LLM_CLIENT = OpenAI(api_key="YOUR_API_KEY")  # TODO: Set your API key or switch client above
+LLM_ID = "gpt-4o"                            # TODO: Change to your chosen model (name it as per the LLM provider's documentation)
+
 LLM_TEMPERATURE = 0.5  # Controls the randomness of the output (the lower, the more deterministic)
 LLM_SEED = 42  # Seed for reproducibility (optional, but recommended)
 LLM_MAX_TOKENS = 3000  # Maximum number of tokens in the generated model (adjust based on your needs)
@@ -28,7 +42,8 @@ LLM_TOP_P = 1.0  # Top-p sampling parameter (1.0 means no filtering)
 TARGET_MODELLING_FRAMEWORK = "CPMpy"
 OUTPUT_FILE = "template_submission.jsonl"
 
-# TODO: Write the main instruction given to the LLM to generate the model.
+# === TODO: Write the main instruction given to the LLM to generate the model. ===
+
 SYSTEM_PROMPT_TEMPLATE = f"""You are an expert in constraint programming.
 Your task is to convert the given natural language problem description into a complete and runnable {TARGET_MODELLING_FRAMEWORK} model.
 The model should be self-contained. 
@@ -38,19 +53,18 @@ The keys in the JSON output should correspond to the decision variables relevant
 Do not include any explanations or introductory text, just the model code between triple backticks.
 For example:
 ```python
-... (the model code here) ...
+# model code here
 ```
 """
 
-# TODO: Write your approach.
+# === MAIN LOGIC, TODO: You can adapt this function to try different prompting strategies ===
 def generate_model_with_llm(problem_id: str, description: str, input_data: str, decision_variables: [str]) -> str:
 
     user_prompt = f"Problem Description:\n{description}\n\n"
-    if decision_variables:
-        user_prompt += (f"The solution should be a JSON object. "
-                        f"The key(s) should strictly be: {', '.join(decision_variables)}.\n\n")
     if input_data:
         user_prompt += f"Input Data:\n{input_data}\n\n"
+    if decision_variables:
+        user_prompt += f"The model must output a JSON with these keys: {', '.join(decision_variables)}.\n\n"
     user_prompt += f"Generate the {TARGET_MODELLING_FRAMEWORK} model."
 
     messages = [
@@ -82,7 +96,7 @@ def generate_model_with_llm(problem_id: str, description: str, input_data: str,
         print(f"  [LLM Call] Error generating model for problem {problem_id}: {type(e).__name__} - {e}")
         return f"# Error generating model for problem {problem_id}: {type(e).__name__} - {e}\n# Please check the LLM configuration and try again."
 
-# --- Main Function (no need to change) ---
+# === MAIN EXECUTION LOOP (No need to change) ===
 def main():
     print(f"Starting model generation script for {TARGET_MODELLING_FRAMEWORK}.")
     print(f"Loading dataset '{GT_DATASET_NAME}'...")