109 lines
2.7 KiB
Python
109 lines
2.7 KiB
Python
import json
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Literal
|
|
|
|
from openai import OpenAI
|
|
from pydantic import BaseModel
|
|
|
|
IN_PATH, OUT_PATH, FAIL_PATH = "data_01/data.json", "data_02/data_processed.json", "data_02/failed.json"
|
|
MAX_TRIES = 2
|
|
WORKERS = 16
|
|
MODEL = "gpt-6-luna"
|
|
|
|
|
|
class Step(BaseModel):
|
|
id: str
|
|
expr: str
|
|
|
|
|
|
class Validation(BaseModel):
|
|
input: dict[str, float]
|
|
process: list[Step]
|
|
output: str
|
|
|
|
|
|
# Structured Outputs(strict)는 dict 스키마를 지원하지 않아서 input을 리스트로 받은 뒤 dict로 변환
|
|
class Var(BaseModel):
|
|
name: str
|
|
value: float
|
|
|
|
|
|
class LLMValidation(BaseModel):
|
|
input: list[Var]
|
|
process: list[Step]
|
|
output: str
|
|
|
|
|
|
class Result(BaseModel):
|
|
status: Literal["success", "unsupported"]
|
|
validation: LLMValidation | None
|
|
|
|
|
|
def evaluate(v: Validation) -> float:
|
|
env = dict(v.input)
|
|
for step in v.process:
|
|
env[step.id] = eval(step.expr, {"__builtins__": {}}, env)
|
|
return env[v.output]
|
|
|
|
|
|
client = OpenAI(
|
|
api_key="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"
|
|
)
|
|
with open("system_prompt.md") as f:
|
|
system_prompt = f.read()
|
|
with open(IN_PATH) as f:
|
|
data = json.load(f)
|
|
|
|
|
|
def convert(item: dict) -> Validation | None:
|
|
response = client.responses.parse(
|
|
model=MODEL,
|
|
reasoning={"effort": "none"},
|
|
instructions=system_prompt,
|
|
input=json.dumps(item, ensure_ascii=False),
|
|
text_format=Result,
|
|
)
|
|
result = response.output_parsed
|
|
if result is None or result.validation is None:
|
|
return None
|
|
|
|
raw = result.validation
|
|
v = Validation(
|
|
input={x.name: x.value for x in raw.input},
|
|
process=raw.process,
|
|
output=raw.output,
|
|
)
|
|
try:
|
|
ok = abs(evaluate(v) - item["answer"]) < 1e-6
|
|
except Exception:
|
|
ok = False
|
|
return v if ok else None
|
|
|
|
|
|
def process(item: dict) -> bool:
|
|
for _ in range(MAX_TRIES):
|
|
try:
|
|
v = convert(item)
|
|
except Exception as e:
|
|
print(f"error: {e}")
|
|
continue
|
|
if v:
|
|
item["validation"] = v.model_dump()
|
|
return True
|
|
return False
|
|
|
|
|
|
start = time.time()
|
|
with ThreadPoolExecutor(max_workers=WORKERS) as pool:
|
|
results = list(pool.map(process, data))
|
|
|
|
processed = [item for item, ok in zip(data, results) if ok]
|
|
failed = [item for item, ok in zip(data, results) if not ok]
|
|
print(f"OK {len(processed)} / FAILED {len(failed)} ({time.time() - start:.1f}s)")
|
|
|
|
with open(OUT_PATH, "w") as f:
|
|
json.dump(processed, f, indent=4, ensure_ascii=False)
|
|
with open(FAIL_PATH, "w") as f:
|
|
json.dump(failed, f, indent=4, ensure_ascii=False)
|