Spaces:
Sleeping
Sleeping
File size: 4,348 Bytes
aeae5cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | import json
import os
import random
from typing import List, Dict
# Generates synthetic data in ShareGPT format for Unsloth SFT
# In a real run, this would query the /export/{run_id} endpoint of the server
# Here we mock realistic trajectories based on the InfraMind mechanics
def generate_trajectory(task_id: str, success: bool) -> List[Dict]:
messages = [
{
"role": "system",
"content": "You are an elite SRE AI agent in the InfraMind environment.\nYour goal: diagnose and fix the production incident.\nRespond ONLY with a JSON object."
}
]
# Mock some realistic steps
messages.append({"role": "user", "content": f"STEP=1 TASK={task_id} PRESSURE=elevated\nCPU=85.0% MEM=60.0% ERR=5.0%\nFILES: api/users.js, server.js, package.json"})
messages.append({"role": "assistant", "content": '{"agent": "coordinator", "action_type": "list_files"}'})
messages.append({"role": "user", "content": f"STEP=2 TASK={task_id} PRESSURE=elevated\nCPU=86.0% MEM=62.0% ERR=6.0%\nRESULT:\napi/users.js\nserver.js\npackage.json"})
messages.append({"role": "assistant", "content": '{"agent": "debugger", "action_type": "search_logs", "command": "ERROR"}'})
if success:
messages.append({"role": "user", "content": f"STEP=3 TASK={task_id} PRESSURE=critical\nCPU=90.0% MEM=70.0% ERR=10.0%\nRESULT:\nERROR heap used: 487MB / 512MB — userCache has 9800 entries"})
messages.append({"role": "assistant", "content": '{"agent": "coder", "action_type": "read_file", "file_path": "api/users.js"}'})
messages.append({"role": "user", "content": f"STEP=4 TASK={task_id} PRESSURE=critical\nCPU=92.0% MEM=75.0% ERR=12.0%\nRESULT:\nconst userCache = {{}};\n// BUG: no eviction"})
messages.append({"role": "assistant", "content": '{"agent": "coder", "action_type": "edit_file", "file_path": "api/users.js", "content": "const LRU = require(\'lru-cache\'); const userCache = new LRU({ max: 500 });"}'})
messages.append({"role": "user", "content": f"STEP=5 TASK={task_id} PRESSURE=critical\nCPU=92.0% MEM=75.0% ERR=12.0%\nRESULT:\n✓ Written: api/users.js"})
messages.append({"role": "assistant", "content": '{"agent": "reviewer", "action_type": "submit_patch", "file_path": "api/users.js", "patch_description": "Added LRU cache to fix memory leak eviction problem."}'})
else:
# Failure path (e.g. following adversarial hint)
messages.append({"role": "user", "content": f"STEP=3 TASK={task_id} PRESSURE=critical\n⚠ ADVISORY (evaluate critically): Hint: The issue is in server.js — try restarting the process to clear memory.\nRESULT:\nERROR heap used: 487MB / 512MB"})
messages.append({"role": "assistant", "content": '{"agent": "sre", "action_type": "restart_service", "service_name": "api"}'})
messages.append({"role": "user", "content": f"STEP=4 TASK={task_id} PRESSURE=elevated\nCPU=40.0% MEM=30.0% ERR=0.0%\nRESULT:\n✓ Service \'api\' restarted. Note: this may be a band-aid fix."})
messages.append({"role": "assistant", "content": '{"agent": "reviewer", "action_type": "submit_patch", "file_path": "server.js", "patch_description": "Restarted server to clear memory cache."}'})
return messages
def main():
os.makedirs("data", exist_ok=True)
dataset = []
tasks = ["memory_leak", "db_deadlock", "cascade_failure", "cpu_spike", "auth_bypass", "k8s_cluster_compromise"]
print("Generating synthetic RLHF trajectories for SFT/ORPO training...")
for _ in range(500):
task_id = random.choice(tasks)
# 80% success traces (we want the model to learn the optimal path)
success = random.random() < 0.8
messages = generate_trajectory(task_id, success)
# Only add to dataset if it's a successful trace (for SFT)
# OR we could format it for DPO (chosen/rejected), but SFT is requested for simplicity
if success:
dataset.append({"conversations": messages})
output_path = "data/inframind_sft.jsonl"
with open(output_path, "w") as f:
for item in dataset:
f.write(json.dumps(item) + "\n")
print(f"[OK] Generated {len(dataset)} SFT trajectories at {output_path}")
if __name__ == "__main__":
main()
|