Update README.md

README.md

@@ -1,1147 +1,234 @@
 ---
 license: apache-2.0
 
+---
+---
+license: apache-2.0
 ---
 
-## Overview of the Main Menu
-
-The `main_menu.py` script is the primary entry point for choosing and executing one of three tasks:
-1. **Training the LLM and World Model**: `train_llm_world`
-2. **Training the Search Agent**: `train_agent`
-3. **Testing the Tree of Thought Search Agent**: `test_agent`
-
-Each task has unique functionalities and configurations. This script uses command-line arguments to specify the desired task and additional options, giving users the ability to tailor the execution according to their needs.
-
-### Running the Main Menu
-
-To run the main menu, use the following command in the terminal:
-```bash
-python main_menu.py --task <task_name> [additional arguments]
-```
-
-Replace `<task_name>` with one of the following:
-- `train_llm_world` - Train the LLM (Language Model) and World Model.
-- `train_agent` - Train the Search Agent with an interactive Twisted-based process.
-- `test_agent` - Test the Tree of Thought Search Agent, with the option of an interactive session or a single query.
-
-### General Arguments
-
-The script supports a set of command-line arguments to customize each task. Here’s an overview of all possible arguments:
+# Model Card for LightBulb
 
-| Argument               | Required | Description                                                                                                                                           | Default                  |
-|------------------------|----------|-------------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------|
-| `--task`               | Yes      | Specifies the task to run. Choose from `train_llm_world`, `train_agent`, or `test_agent`.                                                             | None                     |
-| `--model_name`         | No       | Pretrained model name for LLM. Options include `gpt2`, `bert`, etc., or a custom model path.                                                          | `gpt2`                   |
-| `--dataset_name`       | No       | Name of the dataset from Hugging Face Datasets for training the LLM and World Model (e.g., `wikitext`).                                               | `wikitext`               |
-| `--dataset_config`     | No       | Dataset configuration name for specifying different versions or configurations of the dataset.                                                        | `wikitext-2-raw-v1`      |
-| `--batch_size`         | No       | Number of samples processed in a single forward/backward pass. Increasing the batch size can speed up training but requires more memory.              | `4`                      |
-| `--num_epochs`         | No       | Number of times to iterate over the training dataset during model training. More epochs generally improve learning but can lead to overfitting.       | `3`                      |
-| `--max_length`         | No       | Maximum sequence length for training/inference. Truncates or pads sequences to this length to maintain consistency in training.                      | `128`                    |
-| `--mode`               | No       | Specifies the mode for the LLM and World Model. Use `train` for training and `inference` for generating responses.                                    | `train`                  |
-| `--query`              | No       | Query input for `test_agent` when running a single query instead of an interactive session.                                                           | `''` (empty)             |
+## Overview
 
-## Task Details
+**LightBulb** is an advanced framework designed to train and utilize language models and autonomous web search agents. It integrates hierarchical reinforcement learning, Monte Carlo Tree Search (MCTS), and Tree of Thought (ToT) architectures to enable sophisticated reasoning and decision-making capabilities. The framework supports both training and inference for language models, web search agents, and comprehensive world models.
 
-### 1. Training the LLM and World Model (`train_llm_world`)
+## Installation
 
-This task trains the LLM and the World Model using a chosen dataset from Hugging Face. Training includes adjusting model weights through epochs and creating a model capable of handling long sequences and complex reasoning tasks. 
+To install the necessary dependencies, run:
 
-#### Example Usage
 ```bash
-python main_menu.py --task train_llm_world --model_name gpt2 --dataset_name wikitext --num_epochs 5 --batch_size 8 --max_length 256
+pip install huggingface_hub torch transformers datasets argparse
 ```
 
-#### Arguments Specific to `train_llm_world`
-- **`--model_name`**: Name of the pretrained model to use for language model training. You can specify a model name (like `gpt2`, `bert`, etc.) or a path to a custom model. This argument affects the model architecture and tokenization style.
-  
-- **`--dataset_name`**: Specifies the dataset from Hugging Face’s Datasets library to train the model. Options include `wikitext`, `imdb`, `squad`, etc. You can also use a custom dataset by specifying its path.
-
-- **`--dataset_config`**: Defines the configuration of the dataset, which may be different versions or variations of the dataset. For example, `wikitext` includes configurations such as `wikitext-2-raw-v1`. The configuration will affect the format and content of the data.
-
-- **`--batch_size`**: The number of samples per batch. A larger batch size requires more memory but can improve training speed. You might need to reduce the batch size if memory is limited.
-
-- **`--num_epochs`**: The number of complete passes through the training dataset. More epochs can improve the model’s ability to learn but may lead to overfitting if too high.
-
-- **`--max_length`**: Limits the maximum length of the input sequence. Truncated sequences will be cut off, and shorter sequences will be padded. This affects both training and inference.
-
-- **`--mode`**: Defines the task to be performed. Choose `train` to start training the model. If set to `inference`, the model generates text based on the input.
-
-### 2. Training the Search Agent (`train_agent`)
-Here's a detailed breakdown of your search agent, covering training, inference, and the functionality of each component. This overview will also highlight how the agent saves LLM training data, its modular structure, and the role of each module.
-
----
-
-## Overview of the AutonomousWebAgent
-
-The `AutonomousWebAgent` is a sophisticated, multi-component search and retrieval agent designed to navigate the web, gather relevant content, and perform summarization and generation based on user queries. This agent integrates reinforcement learning (RL), Monte Carlo Tree Search (MCTS), a Retrieva-Augmented Generation (RAG) Summarizer, and a Hierarchical Reinforcement Learning (HRL) architecture to select, execute, and optimize its actions based on past experiences.
-
-### Key Components
-
-1. **Prioritized Experience Replay**:
-   - The agent uses a `PrioritizedReplayMemory` and a `SumTree` to prioritize and store experiences (transitions between states).
-   - The `SumTree` structure maintains a binary tree where each parent node's value is the sum of its children, helping to efficiently store, update, and retrieve experiences based on priority.
-   - These experiences are critical in training both high-level (manager) and low-level (worker) components through prioritized sampling during replay, allowing the model to focus on more significant transitions.
+## Getting Started
 
-2. **Hierarchical Reinforcement Learning (HRL)**:
-   - HRL is employed to allow a **Manager** (high-level) model to select options, which are then executed by a **Worker** (low-level) model. The `ManagerModel` selects tasks (such as searching, summarizing, or generating), while the `WorkerModel` determines specific actions to take.
-   - The manager and worker use LSTM networks with fully connected layers, and each has its own replay memory and optimization process.
-   - The Manager focuses on broad decisions and options, while the Worker operates on specific actions, enabling a layered approach to decision-making.
+### Download the Repository
 
-3. **RAGSummarizer**:
-   - The `RAGSummarizer` leverages a pre-trained language model (e.g., GPT-2) for summarizing, and a SentenceTransformer for embedding-based retrieval. This module breaks down the input content into chunks, retrieves relevant sections based on cosine similarity with the query, and generates a coherent summary.
-   - Additionally, it implements a Least Recently Used (LRU) cache to avoid redundant computation and enhance efficiency, along with persistent storage for cache data.
-   - Summarized results are stored, and this module contributes directly to the generation of LLM training data.
+Use the `huggingface_hub` to download the repository:
 
-4. **WorldModel**:
-   - This module encapsulates an LSTM architecture with linear layers and a `value_head` to estimate state values, allowing the agent to anticipate the long-term value of its actions.
-   - It is utilized in the HRL architecture, specifically by the Worker for evaluating actions and by the Manager in long-term decision-making.
+```python
+from huggingface_hub import snapshot_download
 
-5. **Knowledge Base**:
-   - The knowledge base acts as a repository for collected data, maintaining embeddings for efficient search and retrieval.
-   - It supports saving and loading document embeddings, so the agent can retrieve relevant information for new queries from previously collected knowledge.
-   - Adding and retrieving from the knowledge base enriches the agent’s context and allows it to store and use information from past experiences to inform current tasks.
+# Download the repository
+repo_path = snapshot_download("RobbiePasquale/lightbulb")
 
-6. **Monte Carlo Tree Search (MCTS)**:
-   - The MCTS component guides the agent through complex decision trees to determine the most promising paths for query refinement.
-   - Nodes in the tree represent states (possible query refinements), and child nodes represent possible expansions (e.g., related query variations).
-   - MCTS utilizes a `select`, `expand`, `simulate`, and `backpropagate` strategy to iteratively refine queries, scoring them based on relevance and other metrics to converge on optimal searches.
-   - It also integrates RL by backpropagating rewards based on the ranking score from retrieved results.
-
-7. **Ranking Model**:
-   - The ranking model, built with a neural network and the `SentenceTransformer`, ranks search results based on various features such as cosine similarity with the query, content length, keyword overlap, and domain authority.
-   - This model assigns scores to results, which are then used to guide the MCTS process by enhancing the combined reward with ranking scores.
-
-8. **Tree of Thought (ToT) Search**:
-   - This module enhances the agent's capability to generate a series of interconnected thoughts, exploring different perspectives or angles on a given query.
-   - `ToTNode` and `ToTSearch` classes enable the agent to generate thoughts, evaluate them, and navigate through them as a tree, considering various potential paths to best answer the query.
-   - It combines MCTS and RAG to synthesize responses based on the generated thought paths.
-
-
-### Training Process
-
-The training process for the agent involves episodic learning, where it interacts with various queries from a predefined list. Each query initiates an episode, and the agent performs actions based on its learned policy:
-
-1. **Search and Summarization**:
-   - The agent performs search operations, gathering relevant content from online sources using the MCTS and Ranking Model for prioritization.
-   - Summarization is then carried out on the retrieved content, with relevant information stored in the LLM training data.
-
-2. **Knowledge Base and LLM Training Data Storage**:
-   - Throughout the training process, the agent stores retrieved documents, query results, and summaries in its knowledge base and saves training data for future LLM fine-tuning.
-   - The data is saved in JSONL format and includes metadata such as query terms, source links, and summaries, making it valuable for training language models.
-
-3. **Experience Replay**:
-   - Both the manager and worker models engage in prioritized experience replay, sampling from the stored experiences in the SumTree based on TD-errors.
-   - Replay is essential for reinforcing successful transitions and updating the models' policies over time.
-
-4. **Reward Calculation and Backpropagation**:
-   - Rewards are calculated based on ranking scores, cosine similarity with the query, and other custom factors (e.g., query complexity, state length).
-   - These rewards are backpropagated through the MCTS and used to update the models' decision-making processes, ensuring continuous learning and adaptation.
-
-### Inference Process
-
-During inference:
-- The agent accepts a query, and the Manager model selects a high-level action based on its policy (e.g., search, summarize, or generate).
-- Once an option is chosen, the Worker model executes the corresponding low-level actions. For example, in a search operation, it leverages MCTS to refine the query, retrieves relevant web content, and processes it with the RAGSummarizer.
-- Each inference step is augmented by the agent's existing knowledge base, enabling it to produce more informed and contextually rich responses. Additionally, if Tree of Thought (ToT) is employed, the agent synthesizes a coherent and comprehensive answer based on the thought path.
+print(f"Repository downloaded to: {repo_path}")
+```
 
-### Model Saving
+## Main Features
 
-The agent incorporates a series of save functions to preserve the models:
-- `save_worker_model` and `save_manager_model` functions save the worker and manager models independently.
-- The `save` method preserves the overall state of the agent, which includes its knowledge base, replay memories, and models. This facilitates model reusability and persistent storage, enabling the agent to resume from saved states during training or deployment.
+LightBulb provides six primary functionalities, each accessible via the `main_menu.py` script using command-line arguments.
 
---- 
+### 1. Train a Web Search Agent
 
-This modular setup enhances flexibility, allowing the agent to dynamically adjust its behavior based on rewards from RL, improvements from experience replay, and efficient decision-making through MCTS. Additionally, by saving LLM training data, it becomes highly reusable for further fine-tuning, offering the opportunity to build specialized, data-driven language models optimized for specific domains or tasks.
-This task uses Twisted to train the Autonomous Web Agent by interacting with various queries in a simulated or real environment. It collects rewards based on how well the agent navigates and summarizes web content or performs other tasks.
+**Description:**  
+Trains an autonomous web search agent that navigates the web, gathers relevant content, and learns to summarize and generate responses based on user queries.
 
-#### Example Usage
+**Usage:**
 ```bash
 python main_menu.py --task train_agent
 ```
 
-#### Process Details
-- **Training**: During training, the agent will automatically sample a list of predefined queries, explore web pages, and use reinforcement learning to maximize its reward based on its actions. The training log provides insights into each episode's reward and the agent’s progress.
-  
-- **Logging**: Logs are recorded to `agent_training.log` and provide information about each episode, such as the query, the total reward, and the episode duration. Errors are logged, and if an episode times out, a negative reward is given.
+**Key Components:**
+- **Hierarchical Reinforcement Learning (HRL):** Manages high-level (Manager) and low-level (Worker) decision-making.
+- **Monte Carlo Tree Search (MCTS):** Guides the agent through complex decision trees.
+- **RAGSummarizer:** Summarizes retrieved web content.
+- **Knowledge Base:** Stores and retrieves information to inform future queries.
 
-### 3. Testing the Tree of Thought Search Agent (`test_agent`)
+### 2. Use a Web Search Agent (Inference)
 
-This task lets you test the Tree of Thought Search Agent either in an interactive mode or by specifying a single query. In interactive mode, the user can repeatedly enter queries, and the agent will process them sequentially, producing responses based on the Tree of Thought architecture.
+**Description:**  
+Utilizes the trained web search agent to process queries, perform web searches, and generate summarized responses.
 
-#### Example Usage
-Interactive Mode:
+**Usage:**
 ```bash
 python main_menu.py --task test_agent
 ```
 
-Single Query Mode:
-```bash
-python main_menu.py --task test_agent --query "What are the impacts of renewable energy on global sustainability?"
-```
-
-#### Arguments Specific to `test_agent`
-- **`--query`**: If provided, the agent will process this specific query and return a response. This is ideal for quick, one-off tests or evaluations. If not provided, the program will start an interactive session where you can repeatedly input queries and view the agent's response.
-
-#### Interactive Mode Details
-- **Input**: In interactive mode, enter a query and press Enter. The agent will respond based on its training and the Tree of Thought methodology, traversing different thought paths to generate a response.
-  
-- **Exiting**: To exit the interactive session, type `quit` and press Enter. The agent will then save any new knowledge it has gained and exit the program.
-
-
--------------------------------------------------------------------------------------------------------------------------------------------------
-## World Model
-
-2. **Representation Network**: This module encodes the Transformer output to generate a state representation, reducing dimensionality and making it suitable for further processing.
-3. **Dynamics Network**: This module predicts the next state given a current state and an action. It uses layer normalization and a GELU activation function.
-4. **Prediction Network**: Predicts both the policy logits and value estimates for a given state. It outputs the probabilities of different actions as well as a single scalar value.
-5. **MCTS**: This module performs Monte Carlo Tree Search to evaluate the quality of actions over multiple iterations. It expands nodes based on the policy logits from the Prediction Network and simulates the reward by backpropagating value estimates.
-6. **PPO Agent**: Uses policy and value estimates to calculate PPO loss, which updates the policy while maintaining the constraint on the KL divergence between old and new policies.
-
-The transformer strategically utilises beam search as well as multi token prediction, in order to enrich the encoding from the representation network. 
-
-A generated sequence of tokens is an action, for example if a token is t, then an action is:
-
-a_1= {t1,...,tN}
-
-then a policy is a sequence of actions:
+**Options:**
+- **Interactive Mode:**
+  ```bash
+  python main_menu.py --task test_agent
+  ```
+- **Single Query Mode:**
+  ```bash
+  python main_menu.py --task test_agent --query "Your query here"
+  ```
 
-P_1 = {a_1,...,aN}
+### 3. Train a Language Model
 
-The MCTS and OOPS explores what we are defining as 'thoughts', where a thought is a set of policies:
+**Description:**  
+Trains a Language Model (LLM) and World Model using datasets from Hugging Face, enabling the model to handle complex reasoning and long sequences.
 
-thought_1 = {P1, ... , PN}
-
-The model explores and exploits thoughts, policies, actions, and tokens, and learning happens at each step of granularity. 
-
-
-## Training Details
-
-The model is trained with the following components and techniques:
-
-### Training Procedure
-- **Data Loading**: The data is tokenized and prepared with attention to padding and truncation. Text data is grouped into sequences of fixed length for efficient training.
-- **Optimization**: Training uses an **AdamW** optimizer with **CosineAnnealingLR** scheduler for learning rate adjustments. The **Gradient Scaler** helps prevent overflow when training with mixed precision.
-- **Gradient Accumulation**: Since the model can be computationally heavy, gradients are accumulated over several steps to reduce memory usage.
-- **Loss Functions**: The training process leverages a comprehensive set of custom loss functions:
-
-**1. InfoNCE Loss (Info Noise Contrastive Estimation Loss):**
-Definition: This loss function is used for contrastive learning, encouraging similar samples to be close in the embedding space while pushing dissimilar samples apart.
-
-Formula:
-```
-L_InfoNCE = -log[ exp(sim(z_i, z_j) / τ) / Σ_k exp(sim(z_i, z_k) / τ) ]
-```
-where sim() is the cosine similarity, τ is the temperature parameter, z_i and z_j are paired samples, and the sum in the denominator is over all other samples in the batch.
-
-**2. Covariance Regularization:**
-Definition: This regularization term encourages the learned representations to have uncorrelated dimensions, promoting more diverse and informative embeddings.
-
-Formula:
-```
-L_cov = λ * (Σ_i,j (Cov(i,j)^2 - diag(Cov(i,j))^2))
-```
-where Cov is the covariance matrix of the embeddings, and λ is a regularization coefficient.
-
-**3. Dynamics Performance Loss:**
-Definition: This loss measures the accuracy of predicted next states while also encouraging diverse predictions.
-
-Formula:
-```
-L_dynamics = MSE(true_next_state, predicted_next_state) + λ * Var(predicted_next_state)
-```
-where MSE is the mean squared error, Var is the variance, and λ is a weighting factor.
-
-**4. Thought Consistency Loss:**
-Definition: This loss encourages consistency between true next states and perturbed next states.
-
-Formula:
-```
-L_consistency = MSE(true_next_state, perturbed_next_state)
-```
-
-**5. Policy Value Joint Loss:**
-Definition: This loss combines policy and value losses for reinforcement learning tasks.
-
-Formula:
-```
-L_joint = CrossEntropy(policy_logits, true_policy) + λ * MSE(value_pred, true_value)
+**Usage:**
+```bash
+python main_menu.py --task train_llm_world --model_name gpt2 --dataset_name wikitext --num_epochs 5 --batch_size 8 --max_length 256
 ```
-where λ is a weighting factor balancing policy and value losses.
 
-**6. Action Diversity Reward:**
-Definition: This reward encourages diversity in action embeddings.
+**Key Arguments:**
+- `--model_name`: Pretrained model (e.g., `gpt2`, `bert`).
+- `--dataset_name`: Dataset from Hugging Face (e.g., `wikitext`).
+- `--num_epochs`: Number of training epochs.
+- `--batch_size`: Number of samples per batch.
+- `--max_length`: Maximum sequence length.
 
-Formula:
-```
-R_diversity = λ * Σ_i,j (cos_sim(a_i, a_j)^2)
-```
-where cos_sim is the cosine similarity between action embeddings, and λ is a scaling factor.
+### 4. Inference Using Language Model with Multi-Token Prediction, Beam Search, and MCTS
 
-**7. Expected Thought Value Loss:**
-Definition: This loss aims to maximize the expected value from Monte Carlo Tree Search.
+**Description:**  
+Generates responses using the trained language model, leveraging multi-token prediction, beam search, and MCTS for enhanced coherence and strategic reasoning.
 
-Formula:
-```
-L_ETV = -mean(mcts_best_values)
+**Usage:**
+```bash
+python main_menu.py --task inference_llm --query "Your query here"
 ```
 
-**8. Exploration Regularization:**
-Definition: This regularization encourages exploration by rewarding less-visited actions.
+**Process:**
+1. **Multi-Token Prediction:** Predicts multiple tokens at each step to improve generation speed.
+2. **Beam Search:** Maintains multiple candidate sequences to ensure diverse and high-quality outputs.
+3. **MCTS Integration:** Uses MCTS to evaluate and select the most promising token sequences based on policy and value estimates.
 
-Formula:
-```
-R_exploration = λ * mean(Σ_a (1 / (visit_count(a) + 1)))
-```
-where λ is a scaling factor.
+### 5. Train a Language World Model
 
-**9. KL Divergence Loss:**
-Definition: This loss measures the difference between old and new policies in policy optimization.
+**Description:**  
+Develops a comprehensive World Model that encapsulates state representations, dynamics, and prediction networks to simulate and predict state transitions within the Tree of Thought framework.
 
-Formula:
-```
-L_KL = KL(new_policy || old_policy) = \sum_{i=1}^{n} old\_policy_i \cdot \log\left(\frac{old\_policy_i}{new\_policy_i}\right)
+**Usage:**
+```bash
+python main_menu.py --task train_world_model --additional_args
 ```
-where KL is the Kullback-Leibler divergence.
-
-L_KL is the KL divergence loss
-old_policy and new_policy are probability distributions
-i represents each possible outcome or action
-n is the total number of possible outcomes or actions
-
-### Evaluation
-After each epoch, the model is evaluated on the validation set, computing the average loss over the dataset. The evaluation function utilizes the same loss functions as training but does not backpropagate, allowing it to be run in inference mode.
 
-### Checkpoints
-At the end of each epoch, the model saves checkpoints of all components, enabling easy resumption or further fine-tuning as needed.
+**Key Components:**
+- **Representation Network:** Encodes Transformer outputs into state representations.
+- **Dynamics Network:** Predicts next states based on current states and actions.
+- **Prediction Network:** Generates policy logits and value estimates.
+- **Action Encoder:** Encodes actions into embeddings for state transitions.
 
----
-
-### World Model Components
-
-The World Model encapsulates components that model state representations, dynamics, predictions, and action encodings. These components interact with the Transformer to simulate and predict state transitions within the Tree of Thought framework.
-
-#### RepresentationNetwork
-
-**Function:**
-Transforms the transformer's output embeddings into a compact state representation suitable for modeling and prediction tasks.
-
-**Mathematical Operation:**
-```
-\[
-\text{State} = \text{LayerNorm}\left(\text{Linear}(d_{\text{model}} \rightarrow d_{\text{state}})\left(\text{Linear}(vocab\_dim \rightarrow d_{\text{model}})(\text{Transformer Output})\right)\right)
-\]
-```
-**Explanation:**
-Sequential linear transformations project high-dimensional embeddings into a lower-dimensional state space, followed by layer normalization for stability.
-
-#### DynamicsNetwork
+### 6. Inference with Language World Model
 
-**Function:**
-Models how the state evolves in response to actions (thoughts) taken by the model.
+**Description:**  
+Utilizes the trained World Model to perform advanced reasoning and generate responses based on structured thought processes and state simulations.
 
-**Mathematical Operation:**
-```
-\[
-\text{Next State} = \text{DynamicsNetwork}(\text{Current State}, \text{Action Embedding})
-\]
+**Usage:**
+```bash
+python main_menu.py --task inference_world_model --query "Your query here"
 ```
 
-**Explanation:**
-Predicts the subsequent state by combining the current state representation with an encoded action, effectively simulating the consequences of actions within the Tree of Thought.
+**Features:**
+- **Tree of Thought (ToT):** Structures reasoning paths hierarchically.
+- **Beam Search with MCTS:** Enhances decision-making by balancing exploration and exploitation.
+- **Integration with Knowledge Base:** Leverages stored information for informed responses.
 
-#### PredictionNetwork
+### 7. Inference with World Model, Tree of Thought, and Multi-Token Beam Search
 
-**Function:**
-Predicts policy logits (action probabilities) and value estimates (state evaluations) based on the current state.
+**Description:**  
+Executes inference using the World Model integrated with ToT and multi-token beam search for highly coherent and contextually rich outputs.
 
-**Mathematical Operation:**
-```
-\[
-(\text{Policy Logits}, \text{Value Estimate}) = \text{PredictionNetwork}(\text{State})
-\]
-```
-**Explanation:**
-- **Policy Logits:** Used to derive action probabilities via softmax.
-- **Value Estimate:** Represents the expected reward or quality of the current state.
-
-#### ActionEncoder
-
-**Function:**
-Encodes discrete actions (thoughts) into continuous embeddings compatible with the DynamicsNetwork.
-
-**Mathematical Operation:**
-```
-\[
-\text{Action Embedding} = \text{ActionEncoder}(\text{Action Index})
-\]
+**Usage:**
+```bash
+python main_menu.py --task advanced_inference --query "Your complex query here"
 ```
-**Explanation:**
-Converts action indices into dense vector representations, facilitating their integration into state transition modeling.
-
----
-
-### Tree of Thought (ToT)
-
-The Tree of Thought provides a structured representation of possible thoughts/actions the model can take, organized hierarchically to enable efficient exploration during reasoning.
-
-#### ThoughtNode
-
-**Function:**
-Represents a node in the Tree of Thought, corresponding to a specific action or thought.
 
-**Attributes:**
+**Process:**
+1. **State Initialization:** Converts input queries into state representations.
+2. **MCTS with Beam Search:** Explores multiple reasoning paths simultaneously.
+3. **Thought Sequence Generation:** Produces a sequence of interconnected thoughts/actions.
+4. **Final Response Generation:** Synthesizes the best thought path into a coherent response.
 
-- `name`: Identifier for the thought/action.
-- `children`: List of child `ThoughtNode` instances representing possible subsequent thoughts/actions.
-- `parent`: Reference to the parent `ThoughtNode`.
+## General Arguments
 
-**Mathematical Representation:**
+| Argument           | Required | Description                                                                                      | Default             |
+|--------------------|----------|--------------------------------------------------------------------------------------------------|---------------------|
+| `--task`           | Yes      | Specifies the task to run (`train_llm_world`, `train_agent`, `test_agent`, etc.).               | None                |
+| `--model_name`     | No       | Pretrained model name for LLM (`gpt2`, `bert`, etc.) or a custom model path.                     | `gpt2`              |
+| `--dataset_name`   | No       | Name of the dataset from Hugging Face for training the LLM and World Model (e.g., `wikitext`).   | `wikitext`          |
+| `--dataset_config` | No       | Configuration name for the dataset.                                                              | `wikitext-2-raw-v1` |
+| `--batch_size`     | No       | Number of samples per batch during training.                                                     | `4`                 |
+| `--num_epochs`     | No       | Number of training epochs.                                                                        | `3`                 |
+| `--max_length`     | No       | Maximum sequence length for training/inference.                                                  | `128`               |
+| `--mode`           | No       | Mode for LLM and World Model (`train`, `inference`).                                            | `train`             |
+| `--query`          | No       | Query input for `test_agent` when running a single query.                                       | `''` (empty)        |
 
-Each `ThoughtNode` can be represented as a tree node in a directed graph:
-```
-\[
-\text{ThoughtNode} = (\text{name}, \{\text{children}\})
-\]
-```
-#### State
+## Requirements
 
-**Function:**
-Represents the current state within the MCTS and Tree of Thought framework.
+- **Python:** 3.7+
+- **Libraries:**
+  - `torch>=1.7.1`
+  - `transformers`
+  - `datasets`
+  - `argparse`
+  - `huggingface_hub`
 
-**Attributes:**
+## Usage Examples
 
-- `representation`: Tensor capturing the current state, shaped as \((\text{batch\_size}, \text{seq\_len}, \text{state\_dim})\).
-- `dynamics_network`: Reference to the `DynamicsNetwork` for state transitions.
-- `action_encoder`: Reference to the `ActionEncoder` for encoding actions.
-- `thought_node`: Reference to the current `ThoughtNode` in the Tree of Thought.
+### Training the Language Model and World Model
 
-**Action Application (`apply_action`):**
-```
-\[
-\text{Next State} = \text{DynamicsNetwork}(\text{Current State}, \text{Action Embedding})
-\]
-\[
-\text{New Representation} = \text{Concat}(\text{Current Representation}, \text{Next State} \rightarrow \text{unsqueeze}(1))
-\]
+```bash
+python main_menu.py --task train_llm_world --model_name gpt2 --dataset_name wikitext --num_epochs 5 --batch_size 8 --max_length 256
 ```
-**Procedure:**
 
-1. **Action Encoding:**
-
-```
-   \[
-   \text{Action Index} = \text{Index of Action}
-   \]
-   \[
-   \text{Action Embedding} = \text{ActionEncoder}(\text{Action Index})
-   \]
- ```  
-2. **State Extraction:**
+### Training the Web Search Agent
 
+```bash
+python main_menu.py --task train_agent
 ```
-   \[
-   \text{Current State} = \text{representation}[:, -1, :]
-   \]
-``` 
-3. **State Transition:**
 
-```
-   \[
-   \text{Next State Representation} = \text{DynamicsNetwork}(\text{Current State}, \text{Action Embedding})
-   \]
-```   
-4. **Representation Update:**
+### Testing the Web Search Agent in Interactive Mode
 
+```bash
+python main_menu.py --task test_agent
 ```
-   \[
-   \text{New Representation} = \text{Concat}(\text{representation}, \text{Next State Representation} \times \text{unsqueeze}(1))
-   \]
-```
-5. **Thought Node Update:**
-   - Navigate to the child `ThoughtNode` corresponding to the applied action.
-
----
-
-### Monte Carlo Tree Search (MCTS)
-
-MCTS is an algorithm used to make optimal decisions by traversing the Tree of Thought, balancing exploration of new actions and exploitation of known rewarding actions using statistical methods.
 
-#### MCTSNode
-
-**Function:**
-Represents a node in the MCTS search tree, encapsulating a specific state in the search process.
-
-**Attributes:**
-
-- `state`: Current state represented by a `State` instance.
-- `parent`: Reference to the parent `MCTSNode`.
-- `action`: Action taken to reach this node.
-- `children`: Dictionary mapping actions to child `MCTSNode` instances.
-- `visit_count`: Number of times this node has been visited.
-- `value_sum`: Cumulative value obtained from simulations passing through this node.
-- `prior`: Prior probability of selecting this action, derived from policy logits.
-- `entropy`: Entropy of the policy distribution at this node.
-- `variance`: Variance of the policy distribution at this node.
-
-**Mathematical Representation:**
-
-Each `MCTSNode` can be considered as:
-```
-\[
-\text{MCTSNode} = (\text{state}, \text{parent}, \text{action}, \{\text{children}\}, \text{visit\_count}, \text{value\_sum}, \text{prior}, \text{entropy}, \text{variance})
-\]
-```
-#### MCTS Algorithm
-
-The `MCTS` class implements the Monte Carlo Tree Search algorithm tailored to the LightBulb model's architecture.
-
-**Key Steps:**
-
-1. **Initialization:**
-   - Initialize with the `PredictionNetwork`, `DynamicsNetwork`, `ActionEncoder`, number of iterations (`num_iterations`), exploration constant (`exploration_constant`), beam size (`beam_size`), and number of tokens to predict (`n_tokens_predict`).
-
-2. **Search with Beam (`search_with_beam`):**
-   - **Objective:** Explore the Tree of Thought using beam search augmented with MCTS principles.
-   - **Procedure:**
-     1. **Root Node Evaluation and Backpropagation:**
-        - Evaluate the root node to obtain policy logits and value estimates.
-        - Backpropagate the value estimate to update visit counts and value sums.
-     2. **Beam Initialization:**
-        - Start with a beam containing the root node.
-     3. **Iterative Expansion:**
-        - For each iteration up to `num_iterations`:
-          - **Candidate Collection:**
-            - For each node in the current beam:
-              - If it's a leaf node, evaluate and backpropagate its value.
-              - If it has children, select top `beam_size` actions based on UCB scores.
-              - For each selected action:
-                - Predict a sequence of `n_tokens_predict` actions.
-                - Accumulate scores, entropy, and variance.
-                - Add the candidate sequence to `all_candidates`.
-          - **Beam Pruning:**
-            - Sort all candidates based on a combined score:
-              ```
-              \[
-              \text{Combined Score} = \text{Score} - 0.1 \times \text{Entropy} + 0.05 \times \text{Variance}
-              \]
-              ```
-            - Retain the top `beam_size` candidates for the next iteration.
-     4. **Result Extraction:**
-        - After completing iterations, select the best action sequence from the final beam.
-   
-3. **Evaluation (`evaluate`):**
-   - **Function:** Computes the policy logits and value estimates for a given node's state.
-   - **Procedure:**
-     - Extract the last time step's state representation.
-     - Pass it through the `PredictionNetwork` to obtain policy logits and a value estimate.
-     - Convert logits to probabilities using softmax.
-     - Calculate entropy and variance of the policy distribution.
-     - Expand the node by creating child nodes based on the Tree of Thought and assign priors from policy probabilities.
-   - **Mathematical Operations:**
-     ```
-     \[
-     (\text{Policy Logits}, \text{Value Estimate}) = \text{PredictionNetwork}(\text{State})
-     \]
-     \[
-     P = \text{softmax}(\text{Policy Logits})
-     \]
-     \[
-     \text{Entropy} = -\sum_{i=1}^{V} P_i \log P_i
-     \]
-     \[
-     \text{Variance} = \text{Var}(P)
-     \]
-   ```
-4. **Backpropagation (`backpropagate`):**
-   - **Function:** Updates the `visit_count` and `value_sum` for nodes along the path from the evaluated node back to the root.
-   - **Procedure:**
-```
-     \[
-     \text{For each node in the path:} \\
-     \quad \text{node.visit\_count} \mathrel{+}= 1 \\
-     \quad \text{node.value\_sum} \mathrel{+}= \text{Value Estimate}
-     \]
-```  
-5. **Upper Confidence Bound (UCB) Score (`ucb_score`):**
-   - **Function:** Balances exploration of less-visited nodes and exploitation of high-value nodes.
-   - **Mathematical Operation:**
+### Testing the Web Search Agent with a Single Query
 
+```bash
+python main_menu.py --task test_agent --query "What are the impacts of renewable energy on global sustainability?"
 ```
-     \[
-     \text{UCB Score} = \text{Average Value} + \text{Exploration Term} + \text{Entropy Term} + \text{Variance Term}
-     \]
-     Where:
-     \[
-     \text{Average Value} = \frac{\text{value\_sum}}{\text{visit\_count}}
-     \]
-     \[
-     \text{Exploration Term} = \text{exploration\_constant} \times \text{prior} \times \frac{\sqrt{\text{total\_visits}}}{1 + \text{visit\_count}}
-     \]
-     \[
-     \text{Entropy Term} = -0.1 \times \text{entropy}
-     \]
-     \[
-     \text{Variance Term} = 0.05 \times \text{variance}
-     \]
-``` 
-6. **Best Action Sequence Extraction (`best_action_sequence`):**
-   - **Function:** Extracts the most promising action sequence from the MCTS tree after all iterations.
-   - **Procedure:**
-     - Traverse all possible sequences in the tree.
-     - Score each sequence based on cumulative visit counts, entropy, and variance.
-     - Select the top `beam_size` sequences and return the best one.
-
----
 
-### Beam Search with Multi-Token Prediction
+### Advanced Inference with World Model and Tree of Thought
 
-**Purpose:** Efficiently explores multiple possible token sequences to generate coherent and diverse outputs by predicting multiple tokens at each step.
-
-**Procedure:**
-
-1. **Beam Initialization:**
-
-```
-   - Start with a beam containing the start-of-sequence (BOS) token.
-   \[
-   \text{beam} = \left\{ \left( \text{seq} = [\text{BOS}], \text{score} = 0, \text{cum\_entropy} = 0, \text{cum\_variance} = 0 \right) \right\}
-   \]
-```
-2. **Iterative Expansion:**
-   - For each iteration up to
-```
-\( \frac{\text{max\_length}}{n\_tokens\_predict} \)
-```
-:
-     - For each sequence in the beam:
-       - Predict the next:
+```bash
+python main_menu.py --task advanced_inference --query "Analyze the economic effects of artificial intelligence in the next decade."
 ```
-\( n\_tokens\_predict \)
 
-```
-tokens.
-       - Calculate their probabilities.
-       - Select top-k token sequences based on cumulative scores.
-   
-3. **Beam Pruning:**
-   - After expanding all sequences, retain only the top `beam_size` candidates based on combined scores incorporating log probabilities, entropy, and variance.
+## Citation
 
-4. **Termination:**
-   - Continue until the maximum length is reached or all sequences end with the end-of-sequence (EOS) token.
+If you use LightBulb in your research, please cite the author:
 
-**Mathematical Operations:**
 ```
-\[
-\text{Score} = \sum_{t=1}^{n} \log P(\text{token}_t | \text{tokens}_{<t})
-\]
-\[
-\text{Entropy} = -\sum P_i \log P_i
-\]
-\[
-\text{Variance} = \text{Var}(P)
-\]
+@misc{RobbiePasquale_lightbulb,
+  author       = {Robbie Pasquale},
+  title        = {LightBulb: An Autonomous Web Search and Language Model Framework},
+  year         = {2024},
+  publisher    = {Huggingface},
+  howpublished = {\url{https://huggingface.co/RobbiePasquale/lightbulb}},
+}
 ```
-Where \( P \) is the probability distribution over the vocabulary.
-
-### Upper Confidence Bound (UCB) in MCTS
-
-**Purpose:** Balances exploration of less-visited nodes and exploitation of high-value nodes during tree traversal.
-
-**Mathematical Formulation:**
-\[
-\text{UCB Score} = \text{Average Value} + \text{Exploration Term} + \text{Entropy Term} + \text{Variance Term}
-\]
-Where:
-\[
-\text{Average Value} = \frac{\text{value\_sum}}{\text{visit\_count}}
-\]
-\[
-\text{Exploration Term} = \text{exploration\_constant} \times \text{prior} \times \frac{\sqrt{\text{total\_visits}}}{1 + \text{visit\_count}}
-\]
-\[
-\text{Entropy Term} = -0.1 \times \text{entropy}
-\]
-\[
-\text{Variance Term} = 0.05 \times \text{variance}
-\]
-
-**Explanation:**
-- **Average Value:** Encourages exploitation of nodes with high expected rewards.
-- **Exploration Term:** Encourages visiting nodes with higher uncertainty (low visit counts).
-- **Entropy and Variance Terms:** Modulate preferences based on the policy distribution's entropy and variance, promoting diverse and balanced exploration.
-
-### Entropy and Variance Calculations
-
-**Purpose:** Measure the uncertainty and diversity of the policy distribution, influencing the exploration-exploitation balance.
-
-**Mathematical Formulation:**
-\[
-\text{Entropy} = -\sum_{i=1}^{V} P_i \log P_i
-\]
-\[
-\text{Variance} = \frac{1}{V} \sum_{i=1}^{V} (P_i - \mu)^2
-\]
-Where:
-- \( P_i \): Probability of action \( i \).
-- \( \mu = \frac{1}{V} \sum_{i=1}^{V} P_i \): Mean probability.
-
-**Explanation:**
-- **Entropy:** Quantifies the unpredictability of the distribution. High entropy indicates a more uniform distribution, promoting exploration.
-- **Variance:** Measures the spread of the probabilities. High variance can indicate diverse preferences among actions.
-
----
-
-## Inference Workflow
-
-The inference process in the LightBulb model can operate in two distinct modes:
-
-1. **Without World Model:** Utilizes the Transformer with beam search to generate text directly.
-2. **With World Model and Tree of Thought:** Employs the World Model components alongside MCTS and ToT for generating a sequence of thoughts/actions.
-
-### Inference Modes
-
-#### Language Model
-
-**Procedure:**
-
-1. **Input Processing:**
-   - Tokenize the input query.
-   - Encode tokens into embeddings via the Transformer.
-
-2. **Beam Search Generation:**
-   - Use the Transformer's `generate_with_beam_search` method.
-   - Predict multiple tokens at each step (`n_tokens_predict`).
-   - Maintain a beam of top `beam_size` sequences based on cumulative scores.
-
-3. **Output Decoding:**
-   - Select the best sequence based on scores.
-   - Decode token IDs back into human-readable text.
-
-**Mathematical Operations:**
-
-- **Beam Search Score:**
-  \[
-  \text{Score} = \sum_{t=1}^{n} \log P(\text{token}_t | \text{tokens}_{<t})
-  \]
-
-- **Entropy and Variance:**
-  \[
-  \text{Entropy} = -\sum P_i \log P_i
-  \]
-  \[
-  \text{Variance} = \text{Var}(P)
-  \]
-
-#### World Model and Tree of Thought
-
-**Procedure:**
-
-1. **Input Processing:**
-   - Tokenize the input query.
-   - Encode tokens into embeddings via the Transformer.
-   - Transform the output embeddings into a state representation using the `RepresentationNetwork`.
-
-2. **MCTS Initialization:**
-   - Create a root `MCTSNode` with the initial state and the root `ThoughtNode` from ToT.
-   - Evaluate the root node to obtain policy logits and value estimates.
-   - Backpropagate the evaluation to update visit counts and value sums.
-
-3. **MCTS Iterations with Beam Search:**
-   - For a predefined number of iterations (`num_iterations`):
-     - **Beam Expansion:**
-       - For each node in the current beam:
-         - If it's a leaf node, evaluate and backpropagate its value.
-         - Select top `beam_size` actions based on UCB scores.
-         - For each selected action:
-           - Predict a sequence of `n_tokens_predict` actions.
-           - Accumulate scores, entropy, and variance.
-           - Add the candidate sequence to `all_candidates`.
-     - **Beam Pruning:**
-       - Sort all candidates based on a combined score:
-         \[
-         \text{Combined Score} = \text{Score} - 0.1 \times \text{Entropy} + 0.05 \times \text{Variance}
-         \]
-       - Retain the top `beam_size` candidates for the next iteration.
-
-4. **Output Generation:**
-   - After completing iterations, select the best action sequence from the final beam.
-   - Return the sequence of actions (thoughts) as the output.
 
-**Mathematical Operations:**
 
-- **UCB Score:**
-  \[
-  \text{UCB Score} = \frac{\text{value\_sum}}{\text{visit\_count}} + \text{exploration\_constant} \times \text{prior} \times \frac{\sqrt{\text{total\_visits}}}{1 + \text{visit\_count}} - 0.1 \times \text{entropy} + 0.05 \times \text{variance}
-  \]
+## License
 
-- **Beam Score Adjustment:**
-  \[
-  \text{Combined Score} = \text{Score} - 0.1 \times \text{Entropy} + 0.05 \times \text{Variance}
-  \]
+This project is licensed under the Apache 2.0 License.
 
 ---
 
-## Inference Execution
-
-### Mode: Without World Model
-
-**Step 1: Input Tokenization and Encoding**
-
-1. **Tokenization:**
-   \[
-   \text{input\_ids} = \text{tokenizer.encode(query, return\_tensors='pt')}
-   \]
-   - Shape: \((\text{batch\_size}=1, \text{seq\_len})\)
-
-2. **Encoding via Transformer:**
-   \[
-   \text{transformer\_output} = \text{model\_transformer}(input\_ids, input\_ids)
-   \]
-   - Shape: \((\text{batch\_size}=1, \text{seq\_len}, \text{d\_model})\)
-
-**Step 2: Beam Search Generation**
-
-1. **Beam Search Initialization:**
-   \[
-   \text{beam} = \left\{ \left( \text{seq} = [\text{BOS}], \text{score} = 0, \text{cum\_entropy} = 0, \text{cum\_variance} = 0 \right) \right\}
-   \]
-
-2. **Iterative Expansion:**
-   - For each iteration up to \(\frac{\text{max\_length}}{n\_tokens\_predict}\):
-     - For each sequence in the beam:
-       - If the last token is EOS, retain the sequence.
-       - Else, predict the next \( n\_tokens\_predict \) tokens.
-       - Calculate probabilities, entropy, and variance.
-       - Select top-k tokens for each position based on beam size.
-       - Generate all possible continuations.
-     - **Score Calculation:**
-       \[
-       \text{new\_score} = \text{score} + \sum_{t=1}^{n} \log P(\text{token}_t | \text{tokens}_{<t})
-       \]
-     - **Entropy and Variance Accumulation:**
-       \[
-       \text{new\_entropy} = \text{cum\_entropy} + \sum_{t=1}^{n} \text{Entropy}_t
-       \]
-       \[
-       \text{new\_variance} = \text{cum\_variance} + \sum_{t=1}^{n} \text{Variance}_t
-       \]
-     - **Candidate Aggregation:**
-       - Append new sequences to `all_candidates`.
-   
-3. **Beam Pruning:**
-   - Sort `all_candidates` based on combined scores.
-   - Retain the top `beam_size` candidates for the next iteration.
-
-**Step 3: Output Decoding**
-
-1. **Select Best Sequence:**
-   \[
-   \text{best\_sequence} = \text{beam}[0][0]
-   \]
-   
-2. **Decode Tokens to Text:**
-   \[
-   \text{generated\_text} = \text{tokenizer.decode(best\_sequence, skip\_special\_tokens=True)}
-   \]
-
----
-
-### Mode: With World Model and Tree of Thought
-
-**Step 1: Input Tokenization and Encoding**
-
-1. **Tokenization:**
-   \[
-   \text{input\_ids} = \text{tokenizer.encode(query, return\_tensors='pt')}
-   \]
-   - Shape: \((\text{batch\_size}=1, \text{seq\_len})\)
-
-2. **Encoding via Transformer:**
-   \[
-   \text{transformer\_output} = \text{model\_transformer}(input\_ids, input\_ids)
-   \]
-   - Shape: \((\text{batch\_size}=1, \text{seq\_len}, \text{d\_model})\)
-
-3. **State Representation:**
-   \[
-   \text{initial\_representation} = \text{RepresentationNetwork}(\text{transformer\_output})[:, -1, :].unsqueeze(1)
-   \]
-   - Shape: \((\text{batch\_size}=1, 1, \text{state\_dim})\)
-
-4. **State Initialization:**
-   \[
-   \text{initial\_state} = \text{State}(\text{representation}=\text{initial\_representation}, \text{dynamics\_network}=dynamics\_network, \text{action\_encoder}=action\_encoder, \text{thought\_node}=root\_thought\_node)
-   \]
-
-**Step 2: MCTS Initialization and Root Node Evaluation**
-
-1. **MCTS Instance Creation:**
-   \[
-   \text{mcts} = \text{MCTS}(\text{prediction\_network}, \text{dynamics\_network}, \text{action\_encoder}, \text{num\_iterations}=mcts\_iterations, \text{exploration\_constant}=exploration\_constant)
-   \]
-
-2. **Root Node Creation:**
-   \[
-   \text{root\_node} = \text{MCTSNode}(\text{state}=\text{initial\_state}, \text{thought\_node}=root\_thought\_node)
-   \]
-
-3. **Root Node Evaluation:**
-   \[
-   \text{value\_estimate} = \text{mcts.evaluate}(\text{root\_node})
-   \]
-   
-4. **Backpropagation:**
-   \[
-   \text{mcts.backpropagate}(\text{root\_node}, \text{value\_estimate})
-   \]
-   
-**Step 3: MCTS Iterations with Beam Search**
-
-1. **Beam Initialization:**
-   \[
-   \text{beam} = \left\{ \left( \text{root\_node}, \text{score} = 0, \text{cum\_entropy} = 0, \text{cum\_variance} = 0, \text{action\_sequence} = [] \right) \right\}
-   \]
-
-2. **Iterative Expansion:**
-   - For each iteration up to `num_iterations`:
-     - **Candidate Collection:**
-       - For each node in the current beam:
-         - **Leaf Evaluation:**
-           - If `node.is_leaf()`:
-             \[
-             \text{value\_estimate} = \text{mcts.evaluate}(\text{node})
-             \]
-             \[
-             \text{mcts.backpropagate}(\text{node}, \text{value\_estimate})
-             \]
-         - **Child Selection:**
-           - If `node.children` is not empty:
-             - Calculate `total_visits`:
-               \[
-               \text{total\_visits} = \sum_{\text{child} \in \text{node.children}} \text{child.visit\_count}
-               \]
-             - Select top `beam_size` actions based on UCB scores:
-               \[
-               \text{sorted\_children} = \text{sorted}(\text{node.children.items()}, \text{key}=\lambda \text{item}: \text{item}[1].ucb\_score(\text{total\_visits}, \text{exploration\_constant}), \text{reverse}=True)[:\text{beam\_size}]
-               \]
-         - **Action Sequence Prediction:**
-           - For each selected action:
-             - Initialize `current_node`, `current_sequence`, `current_score`, `current_entropy`, `current_variance`.
-             - **Multi-Token Prediction:**
-               - For each step in `n_tokens_predict`:
-                 - If `current_node.is_leaf()`:
-                   \[
-                   \text{value\_estimate} = \text{mcts.evaluate}(\text{current\_node})
-                   \]
-                   \[
-                   \text{mcts.backpropagate}(\text{current\_node}, \text{value\_estimate})
-                   \]
-                 - If `current_node.children` is empty, break.
-                 - Calculate `total_visits` for the new node.
-                 - Select the action with the highest UCB score:
-                   \[
-                   (\text{next\_action}, \text{next\_node}) = \text{max}(\text{current\_node.children.items()}, \text{key}=\lambda \text{item}: \text{item}[1].ucb\_score(\text{total\_visits}, \text{exploration\_constant}))
-                   \]
-                 - **Score Update:**
-                   \[
-                   \text{current\_score} += \frac{\text{next\_node.value\_sum}}{\text{next\_node.visit\_count}} \quad \text{if} \quad \text{next\_node.visit\_count} > 0 \quad \text{else} \quad 0
-                   \]
-                 - **Entropy and Variance Update:**
-                   \[
-                   \text{current\_entropy} += \text{next\_node.entropy}
-                   \]
-                   \[
-                   \text{current\_variance} += \text{next\_node.variance}
-                   \]
-                 - Append `next_action` to `current_sequence`.
-                 - Update `current_node` to `next_node`.
-             - **Candidate Aggregation:**
-               \[
-               \text{all\_candidates.append}((\text{current\_node}, \text{current\_score}, \text{current\_entropy}, \text{current\_variance}, \text{current\_sequence}))
-               \]
-     
-     - **Beam Pruning:**
-       \[
-       \text{beam} = \text{sorted}(\text{all\_candidates}, \text{key}=\lambda x: x[1] - 0.1 \times x[2] + 0.05 \times x[3], \text{reverse}=True)[:\text{beam\_size}]
-       \]
-   
-3. **Termination:**
-   - Stop early if no candidates remain or all beams have reached terminal nodes.
-
-4. **Result Extraction:**
-   \[
-   \text{best\_sequence} = \text{beam}[0][4]
-   \]
-   - Return `best_sequence` as the generated sequence of actions (thoughts).
-
----
-
-## Integration of Components
-
-The inference process seamlessly integrates multiple components to facilitate advanced reasoning:
-
-1. **Transformer for Sequence Encoding and Generation:**
-   - Processes input sequences and generates embeddings.
-   - Facilitates beam search for direct text generation.
-
-2. **World Model for State Representation and Dynamics:**
-   - `RepresentationNetwork` encodes transformer outputs into state representations.
-   - `DynamicsNetwork` predicts state transitions based on actions.
-   - `PredictionNetwork` provides policy logits and value estimates.
-   - `ActionEncoder` encodes actions into embeddings for state transitions.
-
-3. **Tree of Thought for Structured Reasoning:**
-   - Organizes possible thoughts/actions hierarchically.
-   - Enables systematic exploration of reasoning paths.
-
-4. **Monte Carlo Tree Search for Strategic Exploration:**
-   - Utilizes ToT to explore potential reasoning paths.
-   - Balances exploration and exploitation using UCB scores.
-   - Incorporates beam search with multi-token prediction to handle complex, multi-step actions.
-
-**Workflow Integration:**
-
-- **Input Processing:** Input query is tokenized and encoded via the Transformer.
-- **State Representation:** Encoded inputs are transformed into initial state representations.
-- **MCTS Integration:** MCTS uses the World Model components to explore and evaluate possible thought sequences within the Tree of Thought.
-- **Beam Search:** Multi-token beam search within MCTS ensures diverse and coherent exploration of actions.
-- **Output Generation:** The best sequence of thoughts/actions is extracted and returned as the inference result.
-
----
-
-I am utilising Trees of Thought as a structure of how to structure sets of policies, and sequences of actions. These Tree structures provide the World Model a general thought structure and pattern, similarly to how humans create thought patterns for solving certain problems (e.g. understand, describe, analyse, etc).
-
-Here are some example Trees of Thought:
-graph TD
-    A[Problem-Solving Process] --> B[Problem Identification]
-    A --> C[Problem Analysis]
-    A --> D[Solution Generation]
-    A --> E[Implementation]
-    A --> F[Evaluation and Adjustment]
-    B --> B1[Define the Problem]
-    B --> B2[Identify Stakeholders]
-    B --> B3[Determine Constraints]
-    B --> B4[Recognize Problem Type]
-    B --> B5[Historical Context]
-    C --> C1[Root Cause Analysis]
-    C --> C2[System Mapping]
-    C --> C3[Data Collection]
-    C --> C4[Impact Assessment]
-    C --> C5[Theoretical Framework]
-    D --> D1[Creative Problem Solving]
-    D --> D2[Analytical Approach]
-    D --> D3[Mathematical Computation]
-    D --> D4[Decision Making]
-    E --> E1[Action Planning]
-    E --> E2[Resource Allocation]
-    E --> E3[Change Management]
-    F --> F1[Verification]
-    F --> F2[Performance Metrics]
-    F --> F3[Feedback Loops]
-    F --> F4[Continuous Improvement]
-    C3 --> C3a[Quantitative Data]
-    C3 --> C3b[Qualitative Data]
-    C3 --> C3c[Data Validation]
-    D1 --> D1a[Divergent Thinking]
-    D1 --> D1b[Convergent Thinking]
-    D1 --> D1c[Lateral Thinking]
-    D2 --> D2a[Logical Reasoning]
-    D2 --> D2b[Critical Analysis]
-    D2 --> D2c[Systems Thinking]
-    D3 --> D3a[Basic Operations]
-    D3 --> D3b[Advanced Operations]
-    D3 --> D3c[Computational Methods]
-    D4 --> D4a[Decision Trees]
-    D4 --> D4b[Multi-Criteria Analysis]
-    D4 --> D4c[Probabilistic Reasoning]
-    G[Cross-Cutting Considerations] --> G1[Ethical Framework]
-    G --> G2[Stakeholder Management]
-    G --> G3[Interdisciplinary Connections]
-    G --> G4[Technological Integration]
-    G --> G5[Emotional Intelligence]
-    G --> G6[Collaborative Problem Solving]
-    G1 --> G1a[Value-based Decision Making]
-    G1 --> G1b[Long-term Consequences]
-    G2 --> G2a[Direct Stakeholders]
-    G2 --> G2b[Indirect Stakeholders]
-    G2 --> G2c[Conflicting Interests]
-    G3 --> G3a[Related Fields]
-    G3 --> G3b[Cross-disciplinary Impact]
-    G4 --> G4a[AI-assisted Problem Solving]
-    G4 --> G4b[Data-driven Insights]
-    G4 --> G4c[Digital Collaboration Tools]
-    G5 --> G5a[Self-Awareness]
-    G5 --> G5b[Empathy]
-    G5 --> G5c[Stress Management]
-    G6 --> G6a[Team Dynamics]
-    G6 --> G6b[Communication Strategies]
-    G6 --> G6c[Conflict Resolution]
-    H[Computational Considerations] --> H1[CPU Operations]
-    H --> H2[GPU Parallelization]
-    H --> H3[Floating-Point Precision]
-    I[Order of Operations] --> I1[Parentheses]
-    I --> I2[Exponents]
-    I --> I3[Multiplication and Division]
-    I --> I4[Addition and Subtraction]
-    J[Critical Thinking] --> J1[Assumptions Questioning]
-    J --> J2[Bias Recognition]
-    K[Future Perspective] --> K1[Short-term Projections]
-    K --> K2[Long-term Scenarios]
-    K --> K3[Potential Impacts]
-    L[Learning and Adaptation] --> L1[Reflective Practice]
-    L --> L2[Knowledge Transfer]
-    L --> L3[Adaptive Problem Solving]
-
-
-graph TD
-    A[Meta-Cognitive Strategies] --> B[Creative Problem Solving]
-    A --> C[Systems Thinking]
-    A --> D[Decision Making]
-    A --> E[Emotional Intelligence]
-    A --> F[Collaborative Problem Solving]
-    B --> B1[Divergent Thinking]
-    B --> B2[Convergent Thinking]
-    B --> B3[Lateral Thinking]
-    C --> C1[Holistic Perspective]
-    C --> C2[Feedback Loops]
-    C --> C3[Emergent Properties]
-    D --> D1[Decision Trees]
-    D --> D2[Multi-Criteria Decision Analysis]
-    D --> D3[Probabilistic Reasoning]
-    E --> E1[Self-Awareness]
-    E --> E2[Empathy]
-    E --> E3[Stress Management]
-    F --> F1[Team Dynamics]
-    F --> F2[Communication Strategies]
-    F --> F3[Conflict Resolution]
-    G[Learning and Adaptation]
-    A --> G
-    G --> G1[Reflective Practice]
-    G --> G2[Knowledge Transfer]
-    G --> G3[Adaptive Problem Solving]
-    H[Ethical Framework]
-    A --> H
-    H --> H1[Value-based Decision Making]
-    H --> H2[Stakeholder Analysis]
-    H --> H3[Long-term Consequences]
-    I[Technological Integration]
-    A --> I
-    I --> I1[AI-assisted Problem Solving]
-    I --> I2[Data-driven Insights]
-    I --> I3[Digital Collaboration Tools]
-
-
-## Requirements
-
-This code requires:
-- Python 3.7+
-- `torch>=1.7.1`
-- `transformers`
-- `datasets`
-- `argparse`
-
-
-## Citation
-
-If you use this model in your research, please cite the author.
+For more detailed information on each component and advanced configurations, please refer to the [documentation](https://huggingface.co/RobbiePasquale/lightbulb).