wi-lab
/

lwm

@@ -65,12 +65,12 @@ print("All required functions imported successfully.")
 ---
-### 2. **Load and Tokenize the DeepMIMO Dataset**
-Before loading the LWM model, you need to load the DeepMIMO dataset and select specific scenarios for tokenization. Below is a list of available scenarios and their links for more information:
-| **Scenario**  | **City**      | **Link to DeepMIMO Page**                                                                                     |
-|---------------|---------------|----------------------------------------------------------------------------------------------------------------|
 | Dataset 0    | Denver        | [DeepMIMO City Scenario 18](https://www.deepmimo.net/scenarios/deepmimo-city-scenario18/)                      |
 | Dataset 1    | Indianapolis  | [DeepMIMO City Scenario 15](https://www.deepmimo.net/scenarios/deepmimo-city-scenario15/)                      |
 | Dataset 2    | Oklahoma      | [DeepMIMO City Scenario 19](https://www.deepmimo.net/scenarios/deepmimo-city-scenario19/)                      |
@@ -84,32 +84,46 @@ Before loading the LWM model, you need to load the DeepMIMO dataset and select s
 - **Subcarriers**: 32
 - **Paths**: 20
-#### **Load and Tokenize Code**:
-Select and load specific scenarios by adjusting the `scenario_idxs`. In the example below, we select the first two scenarios and tokenize the data.
 ```python
-# Step 5: Load dataset and select specific scenarios
-print("Loading and tokenizing DeepMIMO dataset...")
 # Load the DeepMIMO dataset
 deepmimo_data = load_DeepMIMO_data()
-# Select scenarios to tokenize
-scenario_idxs = torch.arange(2)  # Adjust the number of scenarios as needed
-# Tokenize the dataset
-preprocessed_chs = tokenizer(deepmimo_data, scenario_idxs, gen_raw=True)
 print("Dataset tokenized successfully.")
 ```
 ---
-### 3. **Load the LWM Model**
-After loading and tokenizing the DeepMIMO dataset, load the pre-trained **LWM** model using the following code:
 ```python
-# Step 6: Load the LWM model (with flexibility for the device)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Loading the LWM model on {device}...")
 model = LWM.from_pretrained(device=device)
@@ -117,12 +131,12 @@ model = LWM.from_pretrained(device=device)
 ---
-### 4. **LWM Inference**
 Once the dataset is tokenized and the model is loaded, generate either **raw channels** or the **inferred LWM embeddings** by choosing the input type.
 ```python
-# Step 7: Generate the dataset for inference
 input_type = ['cls_emb', 'channel_emb', 'raw'][1]  # Modify input type as needed
 dataset = dataset_gen(preprocessed_chs, input_type, model)
 ```
@@ -141,7 +155,7 @@ You can choose between:
 Finally, use the generated dataset for your downstream tasks, such as classification, prediction, or analysis.
 ```python
-# Step 8: Print results
 print(f"Dataset generated with shape: {dataset.shape}")
 print("Inference completed successfully.")
 ```
@@ -151,5 +165,4 @@ print("Inference completed successfully.")
 ## 📋 **Requirements**
 - **Python 3.x**
-- **PyTorch**
-- **Git**

 ---
+### 2. **Load the DeepMIMO Dataset**
+Before tokenizing and processing the data, you need to load the **DeepMIMO** dataset. Below is a list of available datasets and their links for more information:
+| **Dataset**  | **City**      | **Link to DeepMIMO Page**                                                                                     |
+|--------------|---------------|----------------------------------------------------------------------------------------------------------------|
 | Dataset 0    | Denver        | [DeepMIMO City Scenario 18](https://www.deepmimo.net/scenarios/deepmimo-city-scenario18/)                      |
 | Dataset 1    | Indianapolis  | [DeepMIMO City Scenario 15](https://www.deepmimo.net/scenarios/deepmimo-city-scenario15/)                      |
 | Dataset 2    | Oklahoma      | [DeepMIMO City Scenario 19](https://www.deepmimo.net/scenarios/deepmimo-city-scenario19/)                      |
 - **Subcarriers**: 32
 - **Paths**: 20
+#### **Load Data Code**:
+Select and load specific datasets by adjusting the `dataset_idxs`. In the example below, we select the first two datasets.
 ```python
+# Step 5: Load the DeepMIMO dataset
+print("Loading the DeepMIMO dataset...")
 # Load the DeepMIMO dataset
 deepmimo_data = load_DeepMIMO_data()
+# Select datasets to load
+dataset_idxs = torch.arange(2)  # Adjust the number of datasets as needed
+print("DeepMIMO dataset loaded successfully.")
+```
+---
+### 3. **Tokenize the DeepMIMO Dataset**
+After loading the data, tokenize the selected **DeepMIMO** datasets. This step prepares the data for the model to process.
+#### **Tokenization Code**:
+```python
+# Step 6: Tokenize the dataset
+print("Tokenizing the DeepMIMO dataset...")
+# Tokenize the loaded datasets
+preprocessed_chs = tokenizer(deepmimo_data, dataset_idxs, gen_raw=True)
 print("Dataset tokenized successfully.")
 ```
 ---
+### 4. **Load the LWM Model**
+Once the dataset is tokenized, load the pre-trained **LWM** model using the following code:
 ```python
+# Step 7: Load the LWM model (with flexibility for the device)
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Loading the LWM model on {device}...")
 model = LWM.from_pretrained(device=device)
 ---
+### 5. **LWM Inference**
 Once the dataset is tokenized and the model is loaded, generate either **raw channels** or the **inferred LWM embeddings** by choosing the input type.
 ```python
+# Step 8: Generate the dataset for inference
 input_type = ['cls_emb', 'channel_emb', 'raw'][1]  # Modify input type as needed
 dataset = dataset_gen(preprocessed_chs, input_type, model)
 ```
 Finally, use the generated dataset for your downstream tasks, such as classification, prediction, or analysis.
 ```python
+# Step 9: Print results
 print(f"Dataset generated with shape: {dataset.shape}")
 print("Inference completed successfully.")
 ```
 ## 📋 **Requirements**
 - **Python 3.x**
+- **PyTorch**