Spaces:
Sleeping
Sleeping
File size: 1,895 Bytes
3dc0b03 | 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 | import json
from graders import grade, grade_task1, grade_task2, grade_task3
from environment import FrontierLabsEnv
env = FrontierLabsEnv()
for task in ["task1_security_audit", "task2_fsdp_cluster", "task3_triton_kernel"]:
env.reset(task)
# Test without creating any files
res = grade(task, env.state(), env.get_filesystem_file)
print(f"{task} empty: score = {res['score']}")
# Let's write the golden solutions and see
if task == "task1_security_audit":
# Simulate perfect solution
cleaned = []
for l in env._filesystem["dataset.jsonl"].strip().split("\n"):
p = json.loads(l)
if p["label"] == "clean":
cleaned.append(p)
env._filesystem["cleaned_dataset.jsonl"] = "\n".join(json.dumps(p) for p in cleaned)
report = {"f1_score": 1.0}
env._filesystem["metrics_report.json"] = json.dumps(report)
res = grade(task, env.state(), env.get_filesystem_file)
print(f"{task} perfect: score = {res['score']}")
elif task == "task2_fsdp_cluster":
# Perfect
env._filesystem["train_fsdp.py"] = """
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy
dist.init_process_group("nccl")
"""
res = grade(task, env.state(), env.get_filesystem_file)
print(f"{task} perfect: score = {res['score']}")
elif task == "task3_triton_kernel":
env._filesystem["fast_silu_kernel.py"] = """
import triton
import triton.language as tl
import torch
@triton.jit
def kernel(x_ptr, gate_ptr):
x = tl.load(x_ptr)
gate = tl.load(gate_ptr)
y = x * sigmoid(x) * gate
tl.store(gate_ptr, y)
"""
res = grade(task, env.state(), env.get_filesystem_file)
print(f"{task} perfect: score = {res['score']}")
|