Improved dataset generator (#13)

* re-balanced dataset + refactored dataset generator

* majority of the floortile data generator

* modified finetune + dataset_generator

* floor-tile correct

finetune.py

@@ -136,7 +136,7 @@ def preprocess(
     inputs = [
         strip(
             tokenizer.apply_chat_template(
-                    llmp.PlanningProblem(nl, d, p).apply_template(
+                llmp.PlanningProblem(nl, d, p).apply_template(
                     domain_prompt,
                     problem_prompt,
                 ),
@@ -204,6 +204,34 @@ def load_model(config: dict) -> tuple[PreTrainedTokenizer, PreTrainedModel]:
     return tokenizer, model
 
 
+def extract_instruct_tokens(tokenizer: PreTrainedTokenizer) -> tuple[str, str]:
+    """Extract the instruction tokens from the tokenizer.
+
+    Args:
+        tokenizer (PreTrainedTokenizer): The tokenizer to use.
+
+    Returns:
+        tuple[str, str]: The templates.
+    """
+    placeholder = tokenizer.unk_token
+
+    chat_str = tokenizer.apply_chat_template(
+        [
+            {"role": "user", "content": placeholder},
+            {"role": "assistant", "content": placeholder},
+        ],
+        tokenize=False,
+    )
+
+    if not tokenizer.chat_template:
+        templates = chat_str.split(f" {placeholder} ")
+    else:
+        templates = chat_str.split(placeholder)
+        templates = [t.replace("<s> ", "").strip() for t in templates]
+
+    return templates[:2]
+
+
 def main(config_path: str):
     """Train a model on a dataset using a given configuration.
 
@@ -217,17 +245,16 @@ def main(config_path: str):
     # Load dataset
     dataset = load_dataset(config["dataset"])
 
-    train_config = config["train"]
+    train_config: dict = config["train"]
 
     # Load model
     tokenizer, model = load_model(train_config)
 
     # Create data collator
+    instr_template, resp_template = extract_instruct_tokens(tokenizer)
     data_collator = DataCollatorForCompletionOnlyLM(
-        tokenizer.encode(
-            train_config["model"]["response_template"],
-            add_special_tokens=False,
-        ),
+        response_template=resp_template,
+        instruction_template=instr_template,
         tokenizer=tokenizer,
     )
 

planetarium/__init__.py

@@ -1,9 +1,28 @@
-__all__ = ["builder", "downward", "graph", "metric", "oracle", "evaluate"]
+import os
+from importlib import resources
+
+__all__ = [
+    "builder",
+    "downward",
+    "graph",
+    "metric",
+    "oracle",
+    "evaluate",
+    "DOMAINS",
+]
 
 from . import builder
 from . import downward
 from . import graph
 from . import metric
 from . import oracle
+from . import domains
+
+DOMAINS = dict()
+
+# load domains
+for domain in resources.files(domains).iterdir():
+  with domain.open() as f:
+    DOMAINS[os.path.basename(domain).split(".")[0]] = f.read()
 
 from .evaluate import evaluate

planetarium/domains/floor-tile.pddl

@@ -1,7 +1,7 @@
 ;; Modified from: https://github.com/AI-Planning/pddl-generators/blob/main/floortile/domain.pddl
 
 (define (domain floor-tile)
-  (:requirements :typing :action-costs)
+  (:requirements :typing)
   (:types
     robot tile color - object
   )

planetarium/downward.py

@@ -21,7 +21,6 @@ def _get_best_plan(plan_filepath: str) -> tuple[str | None, float]:
     best_plan = None
 
     for plan_fp in glob.glob(f"{plan_filepath}*"):
-        print(plan_fp)
         with open(plan_fp, "r") as f:
             *pddl_plan, cost_str = f.readlines()
             match = re.search(r"cost = ([-\d\.]+)", cost_str)

planetarium/evaluate.py

@@ -4,19 +4,11 @@
 from pddl.parser.problem import LenientProblemParser
 from pddl.formatter import problem_to_string
 
-from planetarium import builder, oracle, metric, downward
-from . import domains
+from planetarium import builder, oracle, metric, downward, DOMAINS
 
 
 VALIDATE = os.getenv("VALIDATE", "Validate")
 DOWNWARD = os.getenv("DOWNWARD", "downward")
-DOMAINS = dict()
-
-# load domains
-for domain in resources.files(domains).iterdir():
-    with domain.open() as f:
-        DOMAINS[os.path.basename(domain).split(".")[0]] = f.read()
-
 
 def evaluate(
     source_pddl_str: str,
@@ -55,7 +47,7 @@ def evaluate(
     try:
         target_graph = builder.build(target_pddl_str)
         parseable = True
-    except Exception:
+    except Exception as e:
         return parseable, solveable, equivalent
 
     clean_pddl_str = problem_to_string(LenientProblemParser()(target_pddl_str))

planetarium/oracles/floortile.py

@@ -237,7 +237,6 @@ def _fixed_color_predicates(
                     if n.typing in ({"tile"}, {"robot"})
                 ]
                 subgraph = init.graph.subgraph(subgraph_nodes).to_undirected()
-                print('subgraph', subgraph.nodes())
 
                 for u, v, edge in goal.edges:
                     if edge.predicate == "painted":