args

app.py

@@ -17,7 +17,7 @@ MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
 
 def main():
-    model_args, data_args, training_args, diffusion_args = get_args()
+    model_args, data_args, training_args, diffusion_args = get_args("args.json")
     tokenizer, model = load_model(model_args, data_args, training_args, diffusion_args, logger)
 
     model.eval()

args.json

@@ -0,0 +1,15 @@
+{
+    "model_name_or_path": "hamishivi/tess2-v0.3",
+    "simplex_value": 5,
+    "num_diffusion_steps": 5000,
+    "num_inference_diffusion_steps": 100,
+    "beta_schedule": "squaredcos_improved_ddpm",
+    "top_p": 0.99,
+    "self_condition": "logits_mean",
+    "self_condition_mix_before_weights": true,
+    "is_causal": false,
+    "mask_padding_in_loss": false,
+    "dataset_name": "c4",
+    "streaming": true,
+    "dataset_config_name": "en"
+}

sdlm/arguments.py

@@ -11,7 +11,7 @@ MODEL_CONFIG_CLASSES = list(MODEL_MAPPING.keys())
 MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
 
 
-def get_args():
+def get_args(filename: str = None):
     parser = HfArgumentParser(
         (
             ModelArguments,
@@ -20,11 +20,11 @@ def get_args():
             DiffusionArguments,
         )
     )
-    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+    if filename is not None:
         # If we pass only one argument to the script and it's the path to a json file,
         # let's parse it to get our arguments.
         model_args, data_args, training_args, diffusion_args = parser.parse_json_file(
-            json_file=os.path.abspath(sys.argv[1])
+            json_file=filename
         )
     else:
         (