""" Spatial Atlas — ML Code Generator Generates complete, self-contained Python scripts for MLE-Bench competitions. Uses strategy templates + LLM to produce runnable ML pipelines. """ import logging import re from llm import LLMClient from mlebench.analyzer import CompetitionAnalysis from mlebench.strategies import get_strategy_template from mlebench.strategies.leaks import leak_prompt_block, match_leak logger = logging.getLogger("spatial-atlas.mlebench.codegen") CODEGEN_SYSTEM_PROMPT = """You are an expert Kaggle grandmaster who writes complete, runnable ML pipeline scripts. Your code must: - Be COMPLETELY self-contained (no human intervention) - Handle edge cases (missing values, unexpected dtypes, empty columns) - Use robust, proven approaches (XGBoost, LightGBM, sklearn, pandas, numpy) - Print progress messages to stdout - Save predictions to the exact submission path specified - Match the submission format EXACTLY (correct columns, correct dtypes) - Include a simple train/validation split for sanity checking - Never crash (wrap risky operations in try/except with fallbacks) - NEVER use display() or show() or plt.show() (headless environment) - On the validation split, print exactly one line of the form VALIDATION_SCORE: where is the competition metric on the held-out fold. This line is machine-parsed to drive score-based refinement. If the competition metric is not directly computable, print the closest proxy (e.g. AUC, accuracy, RMSE) and prefix the line identically.""" CODEGEN_PROMPT = """Generate a complete Python script that solves this Kaggle competition. ## Competition Description {description} ## Available Data Files {file_listing} ## Data Preview {data_preview} ## Analysis - Task type: {task_type} - Metric: {metric} ({metric_direction}) - Target column: {target_column} - Submission format: {submission_format} {leak_block} ## Strategy Template (use as starting point, adapt as needed) {strategy_template} ## Requirements - Read data from: {data_dir}/ - Save submission to: {submission_path} - Use ONLY these libraries: pandas, numpy, sklearn, xgboost, lightgbm, scipy (all pre-installed) - The script must be COMPLETE and RUNNABLE with no edits - Handle missing values explicitly (fillna, dropna, or impute) - Handle categorical columns (label encode or one-hot encode) - Print progress: "Loading data...", "Training model...", "Generating predictions...", "Submission saved." - If target column is unclear, infer it from the submission format - Produce a submission.csv with the exact required format - If the Leak Audit section above fires a real hit, write the leak-derived submission first, then train the baseline anyway as a fallback Generate ONLY the Python code. No markdown fences, no explanation.""" REFINE_PROMPT = """The ML pipeline below ran successfully and produced a valid submission, but we want a stronger score. Propose ONE targeted improvement and return the full updated script. ## Current Script ```python {code} ``` ## Current Validation Score {current_score} (metric: {metric}, higher_is_{metric_direction}) ## Competition Description (first 3000 chars) {description} ## Data Files {file_listing} ## Improvement Menu (pick the most promising ONE for this competition) - Stronger model family (LightGBM -> XGBoost with different params, or add CatBoost) - K-fold cross validation instead of single holdout (report mean score) - Target encoding or more aggressive feature engineering on categoricals - Stacking or blending two model types - Better handling of missing values / outliers - Feature interactions or polynomial features - Class-balanced sampling for imbalanced targets - Hyperparameter tuning via a small grid or Optuna - Leak exploit if the Leak Audit section suggested one ## Requirements - Keep the VALIDATION_SCORE: print line so refinement can continue. - Keep the exact submission path and output format. - Keep the script complete and runnable with no human edits. - Prefer ONE change over many: we will iterate again if this improves things. - Output ONLY the complete Python code. No markdown fences, no explanation. """ FIX_PROMPT = """The ML pipeline script below failed with an error. Fix the script. ## Original Script ```python {code} ``` ## Error {error} ## Stdout Before Error {stdout} ## Competition Description (first 2000 chars) {description} ## Data Files {file_listing} ## Instructions - Fix ONLY the error — don't rewrite the whole script unless necessary - If a file doesn't exist, check the data directory listing for the correct filename - If a column doesn't exist, print available columns first, then adapt - If a library is missing, replace with an available one (pandas, numpy, sklearn, xgboost, lightgbm) - Keep all existing good logic intact - Output ONLY the complete fixed Python script, no explanation""" class MLCodeGenerator: """Generate and fix ML pipeline code for competitions.""" def __init__(self, llm: LLMClient): self.llm = llm async def generate( self, description: str, data_dir: str, file_listing: str, data_preview: str, analysis: CompetitionAnalysis, submission_path: str = "submission.csv", ) -> str: """Generate a complete ML pipeline script.""" strategy_template = get_strategy_template(analysis.strategy) leak_block = leak_prompt_block(description, file_listing) matched = match_leak(description, file_listing) if matched is not None: logger.info( f"Leak hint matched: {matched.slug} ({matched.title}); " "injecting targeted exploit guidance into codegen prompt" ) prompt = CODEGEN_PROMPT.format( description=description[:5000], file_listing=file_listing, data_preview=data_preview[:2000], task_type=analysis.task_type, metric=analysis.metric, metric_direction=analysis.metric_direction, target_column=analysis.target_column, submission_format=analysis.submission_format, leak_block=leak_block, strategy_template=strategy_template, data_dir=data_dir, submission_path=submission_path, ) code = await self.llm.generate( prompt, model_tier="strong", system_prompt=CODEGEN_SYSTEM_PROMPT, max_tokens=8192, temperature=0.1, ) code = self._clean_code(code) logger.info(f"Generated pipeline: {len(code)} chars, {code.count(chr(10))} lines") return code async def refine( self, code: str, current_score: float, metric: str, metric_direction: str, description: str, file_listing: str, ) -> str: """ Propose an improved version of a pipeline that already works. Called after a successful run to drive score-based iteration. The caller is responsible for running the returned code, comparing the new VALIDATION_SCORE against `current_score`, and keeping whichever submission is better. """ prompt = REFINE_PROMPT.format( code=code, current_score=current_score, metric=metric, metric_direction=metric_direction, description=description[:3000], file_listing=file_listing, ) refined = await self.llm.generate( prompt, model_tier="strong", system_prompt=CODEGEN_SYSTEM_PROMPT, max_tokens=8192, temperature=0.3, # slight temperature: we want creative variation ) refined = self._clean_code(refined) logger.info(f"Refined pipeline: {len(refined)} chars") return refined async def fix( self, code: str, error: str, stdout: str, description: str, file_listing: str, ) -> str: """Fix a failed pipeline script based on the error.""" prompt = FIX_PROMPT.format( code=code, error=error[-2000:], stdout=stdout[-1000:], description=description[:2000], file_listing=file_listing, ) fixed_code = await self.llm.generate( prompt, model_tier="strong", system_prompt=CODEGEN_SYSTEM_PROMPT, max_tokens=8192, temperature=0.0, ) fixed_code = self._clean_code(fixed_code) logger.info(f"Fixed pipeline: {len(fixed_code)} chars") return fixed_code def _clean_code(self, code: str) -> str: """Remove markdown fences and clean up generated code.""" code = code.strip() # Remove markdown code fences if code.startswith("```python"): code = code[len("```python"):] elif code.startswith("```"): code = code[3:] if code.endswith("```"): code = code[:-3] code = code.strip() # Ensure there's a newline at the end if not code.endswith("\n"): code += "\n" return code