Update README.md

README.md

@@ -1,5 +1,8 @@
+```markdown
 # 📡 **LWM: Large Wireless Model**
 
+**[🚀 Click here to try the Interactive Demo!](https://huggingface.co/spaces/your-space-name)**
+
 Welcome to the **LWM** (Large Wireless Model) repository! This project hosts a pre-trained model designed to process and extract features from wireless communication datasets, specifically the **DeepMIMO** dataset. Follow the instructions below to clone the repository, load the data, and perform inference with LWM.
 
 ---
@@ -63,34 +66,9 @@ print("All required functions imported successfully.")
 
 ---
 
-### 2. **Load the LWM Model**
-
-Once the repository is cloned, load the pre-trained **LWM** model using the following code:
-
-```python
-# Step 5: 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)
-```
-
----
-
-### 3. **Load the DeepMIMO Dataset**
-
-Load the DeepMIMO dataset using the pre-defined loading function:
-
-```python
-# Step 6: Load dataset (direct call, no module prefix)
-print("Loading DeepMIMO dataset...")
-deepmimo_data = load_DeepMIMO_data()
-```
-
----
-
-### 4. **Tokenize the DeepMIMO Dataset**
+### 2. **Load and Tokenize the DeepMIMO Dataset**
 
-Tokenize the dataset based on specific scenarios from DeepMIMO. Below is a list of available scenarios and their links for more information:
+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**                                                                                     |
 |---------------|---------------|----------------------------------------------------------------------------------------------------------------|
@@ -107,26 +85,45 @@ Tokenize the dataset based on specific scenarios from DeepMIMO. Below is a list
 - **Subcarriers**: 32
 - **Paths**: 20
 
-#### **Tokenization Code**:
-Select and tokenize specific scenarios by adjusting the `scenario_idxs`. In the example below, we select the first two scenarios.
+#### **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 7: Tokenize the dataset
-scenario_idxs = torch.arange(2)  # Adjust the number of scenarios you want
-print("Tokenizing the dataset...")
+# 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.")
 ```
 
-- The dataset will be tokenized according to the selected scenarios and preprocessing configurations.
+---
+
+### 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)
+```
 
 ---
 
-### 5. **LWM Inference**
+### 4. **LWM Inference**
 
-Once the dataset is tokenized, generate either **raw channels** or the **inferred LWM embeddings** by choosing the input type.
+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
+# 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)
 ```
@@ -145,7 +142,7 @@ You can choose between:
 Finally, use the generated dataset for your downstream tasks, such as classification, prediction, or analysis.
 
 ```python
-# Step 9: Print results
+# Step 8: Print results
 print(f"Dataset generated with shape: {dataset.shape}")
 print("Inference completed successfully.")
 ```
@@ -156,3 +153,4 @@ print("Inference completed successfully.")
 
 - **Python 3.x**
 - **PyTorch**
+- **Git**