33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from typing import Any
|
|
|
|
from data.schema import CLASSES, Fact, Rule
|
|
|
|
|
|
def add_rule(record: dict[str, Any], seed: int, held_out: bool = False) -> dict[str, Any]:
|
|
rng = random.Random((seed * 9_999_991) + int(record["scene_id"].split("-")[-1]))
|
|
target = rng.choice(record["objects"])
|
|
color = target["color"]
|
|
color_index = ("red", "blue", "green", "yellow").index(color)
|
|
class_index = (color_index + (1 if held_out else 0)) % len(CLASSES)
|
|
class_name = CLASSES[class_index]
|
|
text_fact = Fact(target["object_id"], "color", color)
|
|
rule = Rule("var_x", "color", color, "class", class_name)
|
|
result = dict(record)
|
|
result.update(
|
|
{
|
|
"rules_text": (
|
|
f"{target['object_id']} is {color}. "
|
|
f"Every {color} object belongs to class {class_name}."
|
|
),
|
|
"text_facts": [text_fact.to_dict()],
|
|
"text_rules": [rule.to_dict()],
|
|
"question": f"Which object belongs to class {class_name}?",
|
|
"answer": target["object_id"],
|
|
"combination": f"{color}->{class_name}",
|
|
}
|
|
)
|
|
return result
|