evolution_engine_mlx.py

"""
代码进化引擎 - MLX 版本
直接使用 MLX-LM 进行本地推理
"""
import json
from datetime import datetime

# 尝试导入 MLX
try:
    from mlx_lm import load, generate
    MLX_AVAILABLE = True
    print("✅ MLX-LM 可用")
except ImportError:
    MLX_AVAILABLE = False
    print("⚠️ MLX-LM 未安装")


class EvolutionEngineMLX:
    def __init__(self, model_name="mlx-community/Qwen2.5-3B-Instruct-4bit"):
        """初始化 MLX 引擎"""
        self.model_name = model_name
        self.available = False
        
        if MLX_AVAILABLE:
            try:
                print(f"🔄 加载模型: {model_name}")
                self.model, self.tokenizer = load(model_name)
                self.available = True
                print("✅ 模型加载成功")
            except Exception as e:
                print(f"❌ 模型加载失败: {e}")
                self.available = False
        else:
            print("⚠️ 使用模拟模式")
    
    def mutate_code(self, code, generation, best_time):
        """生成代码变体"""
        if not self.available:
            return self._mock_mutate(code, generation)
        
        prompt = f"""你是代码优化专家。请优化以下排序函数。

当前代码（第 {generation} 代，耗时 {best_time:.6f}秒）：
```python
{code}
```

优化要求：
1. 保持函数签名不变：def sort_array(arr)
2. 必须返回排序后的数组
3. 尝试改进算法效率（减少时间复杂度）
4. 可以尝试：改变算法逻辑、优化循环、使用更好的数据结构

只返回优化后的完整函数代码，不要解释：
```python"""

        try:
            # 使用 MLX 生成
            response = generate(
                self.model,
                self.tokenizer,
                prompt=prompt,
                max_tokens=500,
                temp=0.7,
                verbose=False
            )
            
            # 提取代码
            code_start = response.find("def sort_array")
            if code_start != -1:
                code_end = response.find("```", code_start)
                if code_end != -1:
                    new_code = response[code_start:code_end].strip()
                else:
                    new_code = response[code_start:].strip()
                
                return new_code
            else:
                return code
                
        except Exception as e:
            print(f"⚠️ 代码生成失败: {e}")
            return code
    
    def _mock_mutate(self, code, generation):
        """模拟模式：返回预设的优化版本"""
        optimizations = [
            # 优化 1：添加提前退出
            """def sort_array(arr):
    n = len(arr)
    for i in range(n):
        swapped = False
        for j in range(n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        if not swapped:
            break
    return arr""",
            # 优化 2：插入排序
            """def sort_array(arr):
    for i in range(1, len(arr)):
        key = arr[i]
        j = i - 1
        while j >= 0 and arr[j] > key:
            arr[j + 1] = arr[j]
            j -= 1
        arr[j + 1] = key
    return arr""",
            # 优化 3：快速排序
            """def sort_array(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[len(arr) // 2]
    left = [x for x in arr if x < pivot]
    middle = [x for x in arr if x == pivot]
    right = [x for x in arr if x > pivot]
    return sort_array(left) + middle + sort_array(right)""",
        ]
        
        idx = min(generation // 5, len(optimizations) - 1)
        return optimizations[idx]
    
    def evolve(self, initial_code, generations=50, variants_per_gen=3):
        """核心进化循环"""
        from code_executor import CodeExecutor
        
        executor = CodeExecutor()
        current_best = initial_code
        best_time = executor.measure_time(current_best)
        
        history = [{
            "generation": 0,
            "code": current_best,
            "time": best_time,
            "timestamp": datetime.now().isoformat()
        }]
        
        mode = "MLX" if self.available else "模拟"
        print(f"\n🧬 开始进化（{mode}模式）...")
        print(f"Generation 0: {best_time:.6f}s (baseline)")
        
        for gen in range(1, generations + 1):
            improved = False
            
            for variant_idx in range(variants_per_gen):
                print(f"\n  生成变体 {variant_idx + 1}/{variants_per_gen}...", end=" ")
                
                variant = self.mutate_code(current_best, gen, best_time)
                variant_time = executor.measure_time(variant)
                
                if variant_time is None:
                    print("❌ 执行失败")
                    continue
                
                print(f"{variant_time:.6f}s", end=" ")
                
                if variant_time < best_time:
                    improvement = ((best_time - variant_time) / best_time) * 100
                    print(f"✅ 提升 {improvement:.1f}%")
                    
                    current_best = variant
                    best_time = variant_time
                    improved = True
                    
                    history.append({
                        "generation": gen,
                        "code": current_best,
                        "time": best_time,
                        "improvement": improvement,
                        "timestamp": datetime.now().isoformat()
                    })
                    
                    break
                else:
                    print("⚪ 无改进")
            
            if not improved:
                print(f"\nGeneration {gen}: 无改进，继续探索...")
            else:
                print(f"\n✨ Generation {gen}: {best_time:.6f}s (best so far)")
        
        print(f"\n🎉 进化完成！")
        print(f"最终性能: {best_time:.6f}s")
        print(f"总提升: {((history[0]['time'] - best_time) / history[0]['time'] * 100):.1f}%")
        
        with open('evolution_history.json', 'w') as f:
            json.dump(history, f, indent=2, ensure_ascii=False)
        
        return {
            "final_code": current_best,
            "final_time": best_time,
            "history": history
        }


if __name__ == "__main__":
    initial_code = """def sort_array(arr):
    n = len(arr)
    for i in range(n):
        for j in range(n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    return arr"""
    
    engine = EvolutionEngineMLX()
    result = engine.evolve(initial_code, generations=5, variants_per_gen=2)
    
    print("\n" + "="*50)
    print("最终代码：")
    print("="*50)
    print(result['final_code'])