Spaces:
Sleeping
Sleeping
| import asyncio | |
| import os | |
| import random | |
| import csv | |
| import json | |
| import statistics | |
| import time | |
| import subprocess | |
| import socket | |
| from datetime import datetime | |
| from typing import List, Dict | |
| # Standardize environment variables BEFORE importing inference | |
| # This handles the case where users set variables with quotes in CMD | |
| for key in ["API_BASE_URL", "API_KEY", "MODEL_NAME"]: | |
| val = os.environ.get(key) | |
| if val: | |
| # Strip literal quotes that CMD sometimes includes | |
| os.environ[key] = val.strip('"').strip("'") | |
| # Set defaults if not provided | |
| if "API_BASE_URL" not in os.environ: | |
| os.environ["API_BASE_URL"] = "http://127.0.0.1:11434/v1" | |
| if "API_KEY" not in os.environ: | |
| os.environ["API_KEY"] = "ollama" | |
| # (MODEL_NAME is handled inside inference.py default, but we can set it here too) | |
| if "MODEL_NAME" not in os.environ: | |
| os.environ["MODEL_NAME"] = "qwen3.5:latest" | |
| # Import the core logic from inference.py | |
| from inference import ( | |
| MaterialForgeEnv, | |
| run_episode, | |
| start_environment_server, | |
| TASKS as ORIGINAL_TASKS, | |
| MODEL_NAME, | |
| SUCCESS_THRESHOLD | |
| ) | |
| from openai import OpenAI | |
| def is_port_open(port): | |
| with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: | |
| return s.connect_ex(('127.0.0.1', port)) == 0 | |
| def start_ollama_server(): | |
| if is_port_open(11434): | |
| print("[INFO] Ollama server is already running on port 11434.") | |
| return None | |
| print("[INFO] Ollama server not detected. Starting 'ollama serve'...") | |
| try: | |
| proc = subprocess.Popen( | |
| ["ollama", "serve"], | |
| stdout=subprocess.DEVNULL, | |
| stderr=subprocess.DEVNULL, | |
| creationflags=subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.DETACHED_PROCESS if os.name == 'nt' else 0 | |
| ) | |
| # Wait up to 30 seconds for Ollama to become ready | |
| start_time = time.time() | |
| while time.time() - start_time < 30: | |
| if is_port_open(11434): | |
| print("[INFO] Ollama server started successfully.") | |
| # Give it an extra second to initialize routes | |
| time.sleep(2) | |
| return proc | |
| time.sleep(1) | |
| print("[WARNING] Ollama server start timed out. It may still be warming up.") | |
| return proc | |
| except Exception as e: | |
| print(f"[ERROR] Failed to start Ollama server: {e}") | |
| return None | |
| async def run_analytical_suite(num_trials: int = 100, mode: str = "llm"): | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| report_file = f"benchmark_{mode}_{timestamp}.csv" | |
| summary_file = f"summary_{mode}_{timestamp}.txt" | |
| ollama_proc = None | |
| if mode == "llm": | |
| ollama_proc = start_ollama_server() | |
| print(f"📊 Starting Analytical Benchmark Task ({mode.upper()} mode)") | |
| print(f"🕒 Timestamp: {timestamp}") | |
| print(f"📝 Logging results to: {report_file}") | |
| api_url = os.environ.get("API_BASE_URL", "http://localhost:11434/v1") | |
| api_key = os.environ.get("API_KEY", "ollama") | |
| # In heuristic mode, we don't need a real client | |
| client = OpenAI(base_url=api_url, api_key=api_key) if mode == "llm" else None | |
| if mode == "heuristic": | |
| print("🛠 Mode: Pure Heuristic (LLM Bypassed)") | |
| else: | |
| print(f"🤖 Model: {os.getenv('MODEL_NAME', MODEL_NAME)}") | |
| # Start the environment server once | |
| server_proc = start_environment_server(port=7860) | |
| env = MaterialForgeEnv(base_url="http://127.0.0.1:7860") | |
| try: | |
| await env.connect() | |
| except Exception as e: | |
| print(f"❌ Connection failed: {e}. Is Ollama or the server running?") | |
| if server_proc: server_proc.terminate() | |
| return | |
| results = [] | |
| successes = 0 | |
| print("\n" + "="*50) | |
| print(f"{'Trial':<8} | {'Seed':<8} | {'Score':<8} | {'Status':<10}") | |
| print("-" * 50) | |
| with open(report_file, mode='w', newline='') as csvfile: | |
| fieldnames = ['trial', 'seed', 'score', 'success', 'timestamp'] | |
| writer = csv.DictWriter(csvfile, fieldnames=fieldnames) | |
| writer.writeheader() | |
| for i in range(1, num_trials + 1): | |
| seed = random.randint(1000, 9999) | |
| # Select a random task type from the original list as a template | |
| template = random.choice(ORIGINAL_TASKS) | |
| task = { | |
| "name": f"analytical-{i}", | |
| "difficulty": template["difficulty"], | |
| "seed": seed | |
| } | |
| try: | |
| # Run the episode without printing every single [STEP] to keep log clean | |
| # We wrap it to suppress standard prints if we wanted to, but let's keep it simple | |
| score = await run_episode(env, client, task) | |
| is_success = score >= SUCCESS_THRESHOLD | |
| results.append(score) | |
| if is_success: | |
| successes += 1 | |
| status = "✅ PASS" if is_success else "❌ FAIL" | |
| print(f"{i:<8} | {seed:<8} | {score:<8.3f} | {status:<10}") | |
| writer.writerow({ | |
| 'trial': i, | |
| 'seed': seed, | |
| 'score': f"{score:.3f}", | |
| 'success': is_success, | |
| 'timestamp': datetime.now().isoformat() | |
| }) | |
| except Exception as e: | |
| print(f"Trial {i} Error: {e}") | |
| # Final Analysis | |
| summary_data = [ | |
| "="*50, | |
| "📈 FINAL ANALYTICAL REPORT", | |
| "="*50, | |
| f"Mode: {mode.upper()}", | |
| f"Model: {os.getenv('MODEL_NAME', MODEL_NAME) if mode == 'llm' else 'Heuristic-Only'}", | |
| f"Total Trials: {num_trials}", | |
| f"Success Rate: {(successes/num_trials)*100:.1f}%", | |
| f"Average Score: {statistics.mean(results):.3f}", | |
| f"Max Score: {max(results):.3f}", | |
| f"Min Score: {min(results):.3f}", | |
| "="*50 | |
| ] | |
| report_text = "\n".join(summary_data) | |
| print("\n" + report_text) | |
| with open(summary_file, 'w') as f: | |
| f.write(report_text) | |
| # Cleanup | |
| await env.close() | |
| if server_proc: | |
| server_proc.terminate() | |
| server_proc.wait() | |
| if ollama_proc: | |
| print("[INFO] Terminating local Ollama server...") | |
| ollama_proc.terminate() | |
| try: | |
| ollama_proc.wait(timeout=5) | |
| except subprocess.TimeoutExpired: | |
| ollama_proc.kill() | |
| if __name__ == "__main__": | |
| import argparse | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--trials", type=int, default=100) | |
| parser.add_argument("--mode", type=str, choices=["llm", "heuristic"], default="llm") | |
| args = parser.parse_args() | |
| asyncio.run(run_analytical_suite(args.trials, mode=args.mode)) | |