Update main_menu.py

main_menu.py

@@ -1,61 +1,118 @@
-# main_menu.py
-
-import argparse
-import sys
-from train_agent import train_agent
-from test_agent import TestAgent, run_test_session
-from lightbulb import main as world_model_main
-
-def parse_main_args():
-    parser = argparse.ArgumentParser(description="Main Menu for Selecting Tasks")
-    parser.add_argument('--task', type=str, choices=['train_llm_world', 'train_agent', 'test_agent'], 
-                        required=True, help='Choose task to execute: train_llm_world, train_agent, test_agent')
-    # Optional arguments for more granular control
-    parser.add_argument('--model_name', type=str, default='gpt2', help='Pretrained model name for LLM')
-    parser.add_argument('--dataset_name', type=str, default='wikitext', help='Dataset name for training')
-    parser.add_argument('--dataset_config', type=str, default='wikitext-2-raw-v1', help='Dataset configuration name')
-    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for training')
-    parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs for training')
-    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length for training')
-    parser.add_argument('--mode', type=str, choices=['train', 'inference'], default='train', help='Train or inference mode for LLM')
-    parser.add_argument('--query', type=str, default='', help='Query for the test_agent')
-    return parser.parse_args()
-
-def main():
-    # Parse arguments for the main function
-    args = parse_main_args()
-
-    # Execute tasks based on user input
-    if args.task == 'train_llm_world':
-        print("Starting LLM and World Model Training...")
-        # Directly call the world model main function
-        sys.argv = ['lightbulb.py', '--mode', args.mode, '--model_name', args.model_name, 
-                    '--dataset_name', args.dataset_name, '--dataset_config', args.dataset_config,
-                    '--batch_size', str(args.batch_size), '--num_epochs', str(args.num_epochs),
-                    '--max_length', str(args.max_length)]
-        world_model_main()
-
-    elif args.task == 'train_agent':
-        print("Starting Agent Training...")
-        # Call the train_agent function from train_agent.py
-        from twisted.internet import reactor, task
-        d = task.deferLater(reactor, 0, train_agent)
-        d.addErrback(lambda failure: print(f"An error occurred: {failure}", exc_info=True))
-        d.addBoth(lambda _: reactor.stop())
-        reactor.run()
-
-    elif args.task == 'test_agent':
-        print("Starting Test Agent...")
-        test_agent = TestAgent()
-        if args.query:
-            # Directly process a single query
-            result = test_agent.process_query(args.query)
-            print("\nAgent's response:")
-            print(result)
-        else:
-            # Run the interactive session
-            reactor.callWhenRunning(run_test_session)
-            reactor.run()
-
-if __name__ == "__main__":
-    main()
+# main_menu.py
+
+import argparse
+import sys
+import os
+from train_agent import train_agent
+from test_agent import TestAgent, run_test_session
+from lightbulb import main as world_model_main
+from lightbulb_inf import main as inference_main  # Import the inference main function
+from twisted.internet import reactor, task
+
+def parse_main_args():
+    parser = argparse.ArgumentParser(description="Main Menu for Selecting Tasks")
+    parser.add_argument('--task', type=str, choices=[
+                        'train_llm_world', 
+                        'train_agent', 
+                        'test_agent', 
+                        'inference_llm', 
+                        'inference_world_model', 
+                        'advanced_inference'
+                    ], 
+                        required=True, 
+                        help='Choose task to execute: train_llm_world, train_agent, test_agent, inference_llm, inference_world_model, advanced_inference')
+    # Optional arguments for more granular control
+    parser.add_argument('--model_name', type=str, default='gpt2', help='Pretrained model name for LLM')
+    parser.add_argument('--dataset_name', type=str, default='wikitext', help='Dataset name for training')
+    parser.add_argument('--dataset_config', type=str, default='wikitext-2-raw-v1', help='Dataset configuration name')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for training')
+    parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs for training')
+    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length for training')
+    parser.add_argument('--mode', type=str, choices=['train', 'inference'], default='train', help='Train or inference mode for LLM')
+    parser.add_argument('--query', type=str, default='', help='Query for the test_agent or inference tasks')
+    # Additional arguments specific to inference can be added here if needed
+    return parser.parse_args()
+
+def main():
+    # Parse arguments for the main function
+    args = parse_main_args()
+
+    # Execute tasks based on user input
+    if args.task == 'train_llm_world':
+        print("Starting LLM and World Model Training...")
+        # Directly call the world model main function with appropriate arguments
+        sys.argv = [
+            'lightbulb.py', 
+            '--mode', args.mode, 
+            '--model_name', args.model_name, 
+            '--dataset_name', args.dataset_name, 
+            '--dataset_config', args.dataset_config,
+            '--batch_size', str(args.batch_size), 
+            '--num_epochs', str(args.num_epochs),
+            '--max_length', str(args.max_length)
+        ]
+        world_model_main()
+
+    elif args.task == 'train_agent':
+        print("Starting Agent Training...")
+        # Call the train_agent function from train_agent.py using Twisted reactor
+        d = task.deferLater(reactor, 0, train_agent)
+        d.addErrback(lambda failure: print(f"An error occurred: {failure}", exc_info=True))
+        d.addBoth(lambda _: reactor.stop())
+        reactor.run()
+
+    elif args.task == 'test_agent':
+        print("Starting Test Agent...")
+        test_agent = TestAgent()
+        if args.query:
+            # Directly process a single query
+            result = test_agent.process_query(args.query)
+            print("\nAgent's response:")
+            print(result)
+        else:
+            # Run the interactive session
+            reactor.callWhenRunning(run_test_session)
+            reactor.run()
+
+    elif args.task in ['inference_llm', 'inference_world_model', 'advanced_inference']:
+        print("Starting Inference Task...")
+        # Prepare the arguments for lightbulb_inf.py based on the selected inference task
+
+        # Map the main_menu task to lightbulb_inf.py's inference_mode
+        inference_mode_map = {
+            'inference_llm': 'without_world_model',
+            'inference_world_model': 'world_model',
+            'advanced_inference': 'world_model_tree_of_thought'
+        }
+
+        selected_inference_mode = inference_mode_map.get(args.task, 'world_model_tree_of_thought')
+
+        # Construct sys.argv for lightbulb_inf.py
+        lightbulb_inf_args = [
+            'lightbulb_inf.py',
+            '--mode', 'inference',
+            '--model_name', args.model_name,
+            '--query', args.query,
+            '--max_length', str(args.max_length),
+            '--inference_mode', selected_inference_mode,
+            '--beam_size', str(getattr(args, 'beam_size', 5)),
+            '--n_tokens_predict', str(getattr(args, 'n_tokens_predict', 3)),
+            '--mcts_iterations', str(getattr(args, 'mcts_iterations', 10)),
+            '--mcts_exploration_constant', str(getattr(args, 'mcts_exploration_constant', 1.414))
+        ]
+
+        # Include additional arguments if they exist
+        if hasattr(args, 'load_model') and args.load_model:
+            lightbulb_inf_args += ['--load_model', args.load_model]
+
+        # Update sys.argv and call the inference main function
+        sys.argv = lightbulb_inf_args
+        inference_main()
+
+    else:
+        print(f"Unknown task: {args.task}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()
+