<ADD> optimize code

app.py

@@ -56,17 +56,10 @@ def do_infer_probs(model, exemplar_attn_kv, exemplar_attn_mask, batched_choices_
 
 @torch.no_grad()
 def process_once(dataset_name, exemplar_str, forward_steps, raw_data):
-    model_name, model_size = "opt", "125m"
-    step_size, momentum = 0.01, 0.9
-
     setup_cpu(seed=seed)
     TaskHandler = load_task(dataset_name)
     task_agent = TaskHandler(prompt_version)
 
-    tokenizer = build_tokenizer(model_name, model_size, padding_side="right")
-    model = build_model(model_name, model_size, False)
-    torch.autograd.set_grad_enabled(False)
-
     processed_data = task_agent.dataset_preprocess(raw_data)
     dataset = TokenizedForMCRightPad(processed_data, tokenizer, task_agent.multiple_choice_promptify)
 
@@ -132,6 +125,13 @@ if __name__ == "__main__":
     prompt_version = "default"
     kv_iter = 10
 
+    model_name, model_size = "opt", "125m"
+    step_size, momentum = 0.01, 0.9
+    setup_cpu(seed=seed)
+    tokenizer = build_tokenizer(model_name, model_size, padding_side="right")
+    model = build_model(model_name, model_size, False)
+    torch.autograd.set_grad_enabled(False)
+
     print(f"Dataset: {dataset_name}")
     task_root = Path("example_sets").joinpath(dataset_name)