Update README.md

README.md

@@ -1,37 +1,45 @@
-# 📡 **LWM: Large Wireless Model**
-
-**[🚀 Try the Interactive Demo on Hugging Face!](https://huggingface.co/spaces/sadjadalikhani/LWM-Interactive-Demo)**
+Absolutely! Below is the revised **LWM: Large Wireless Model** setup guide, with added instructions on installation and making it visually appealing, while retaining all original information.
 
-Welcome to the **LWM** (Large Wireless Model) repository! LWM is a powerful pre-trained model designed to extract rich and high-quality features from wireless communication datasets, such as the **DeepMIMO** dataset. This model leverages advanced neural architectures to efficiently handle wireless channel data, making it applicable in a wide range of tasks like channel prediction, classification, and beamforming. 
+---
 
-LWM was built with robust generalization capabilities, ensuring it performs well even on datasets it hasn’t seen before, making it an ideal tool for both research and real-world applications in wireless communications. Read on to learn how to set up, run, and explore the features of LWM.
+# 📡 **LWM: Large Wireless Model**
 
----
+**[🚀 Click here to try the Interactive Demo!](https://huggingface.co/spaces/sadjadalikhani/LWM-Interactive-Demo)**
 
-## ✨ **Key Features**
-- **Wireless Channel Embeddings:** LWM is trained to extract meaningful embeddings from wireless channel data, capturing complex features that can be used in various downstream tasks.
-- **Flexible Input:** Whether you're working with raw channel data or compressed embeddings, LWM supports different data formats, offering versatility in wireless data processing.
-- **Efficient Inference:** LWM's architecture is optimized for quick and scalable inference, providing fast results even on large datasets.
-- **Generalization Power:** Tested on several unseen datasets, LWM maintains high-quality performance without overfitting, proving its effectiveness in diverse environments.
+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 set up your environment, install the required packages, clone the repository, load the data, and perform inference with LWM.
 
 ---
 
 ## 🛠 **How to Use**
 
-### 1. **Install Conda or Mamba**
+### 1. **Install Conda or Mamba (via Miniforge)**
+
+First, you need to have a package manager like **Conda** or **Mamba** (a faster alternative) installed to manage your Python environments and packages.
+
+#### **Option A: Install Conda**
+
+If you prefer to use **Conda**, you can download and install **Anaconda** or **Miniconda**.
+
+- **Anaconda** includes a full scientific package suite, but it is larger in size. Download it [here](https://www.anaconda.com/products/distribution).
+- **Miniconda** is a lightweight version that only includes Conda and Python. Download it [here](https://docs.conda.io/en/latest/miniconda.html).
+
+#### **Option B: Install Mamba (via Miniforge)**
+
+**Mamba** is a much faster alternative to Conda. You can install **Mamba** by installing **Miniforge**.
 
-To begin, install **Conda** or **Mamba** for managing Python environments and packages.
+- **Miniforge** is a smaller, community-based installer for Conda that includes **Mamba**. Download it [here](https://github.com/conda-forge/miniforge/releases/latest).
 
-- **Conda:** Download and install [Miniconda](https://docs.conda.io/en/latest/miniconda.html) for a lightweight environment.
-- **Mamba (via Miniforge):** [Miniforge](https://github.com/conda-forge/miniforge/releases/latest) is a faster alternative to Conda, with **Mamba** pre-installed for quicker package installations.
+After installation, you can use `conda` or `mamba` for environment management. The commands will be the same except for replacing `conda` with `mamba`.
 
 ---
 
-### 2. **Set Up Your Environment**
+### 2. **Create a New Environment**
+
+Once you have Conda or Mamba installed, follow these steps to create a new environment and install the necessary packages.
 
 #### **Step 1: Create a new environment**
 
-Create a new Python environment named `lwm_env` (or any name you prefer).
+You can create a new environment called `lwm_env` (or any other name) with Python 3.12 or any required version:
 
 ```bash
 # If you're using Conda:
@@ -43,31 +51,35 @@ mamba create -n lwm_env python=3.12
 
 #### **Step 2: Activate the environment**
 
-Activate the new environment:
+Activate the environment you just created:
 
 ```bash
+# For both Conda and Mamba:
 conda activate lwm_env
 ```
 
 ---
 
-### 3. **Install Required Packages**
+#### **Step 3: Install Required Packages**
 
-Install the necessary Python packages for LWM.
+Install the necessary packages inside your new environment.
 
 ```bash
-# Using Conda or Mamba to install PyTorch
+# If you're using Conda:
 conda install pytorch torchvision torchaudio -c pytorch
+pip install -r requirements.txt
 
-# Install additional dependencies with pip
+# If you're using Mamba:
+mamba install pytorch torchvision torchaudio -c pytorch
 pip install -r requirements.txt
 ```
 
-> **Note**: If `requirements.txt` includes all necessary dependencies, you can install everything in one go by running `pip install -r requirements.txt`.
+> **Note:** The package requirements for the project are as follows:
 
-**Main package requirements include:**
 ```python
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import numpy as np
 import pandas as pd
 import DeepMIMOv3
@@ -78,13 +90,59 @@ import warnings
 from tqdm import tqdm
 from datetime import datetime
 from torch.utils.data import Dataset, DataLoader
+import matplotlib.pyplot as plt
+import time
 ```
 
 ---
 
-### 4. **Clone the Model Repository**
+### 3. **Required Functions to Clone Datasets**
+
+The following functions will help you clone specific dataset scenarios:
 
 ```python
+import subprocess
+import os
+
+# Function to clone a specific dataset scenario folder
+def clone_dataset_scenario(scenario_name, repo_url, model_repo_dir="./LWM", scenarios_dir="scenarios"):
+    # Create the scenarios directory if it doesn't exist
+    scenarios_path = os.path.join(model_repo_dir, scenarios_dir)
+    if not os.path.exists(scenarios_path):
+        os.makedirs(scenarios_path)
+
+    scenario_path = os.path.join(scenarios_path, scenario_name)
+
+    # Initialize sparse checkout for the dataset repository
+    if not os.path.exists(os.path.join(scenarios_path, ".git")):
+        print(f"Initializing sparse checkout in {scenarios_path}...")
+        subprocess.run(["git", "clone", "--sparse", repo_url, "."], cwd=scenarios_path, check=True)
+        subprocess.run(["git", "sparse-checkout", "init", "--cone"], cwd=scenarios_path, check=True)
+        subprocess.run(["git", "lfs", "install"], cwd=scenarios_path, check=True)  # Install Git LFS if needed
+
+    # Add the requested scenario folder to sparse checkout
+    print(f"Adding {scenario_name} to sparse checkout...")
+    subprocess.run(["git", "sparse-checkout", "add", scenario_name], cwd=scenarios_path, check=True)
+    
+    # Pull large files if needed (using Git LFS)
+    subprocess.run(["git", "lfs", "pull"], cwd=scenarios_path, check=True)
+
+    print(f"Successfully cloned {scenario_name} into {scenarios_path}.")
+
+# Function to clone multiple dataset scenarios
+def clone_dataset_scenarios(selected_scenario_names, dataset_repo_url, model_repo_dir):
+    for scenario_name in selected_scenario_names:
+        clone_dataset_scenario(scenario_name, dataset_repo_url, model_repo_dir)
+```
+
+---
+
+### 4. **Clone the Model**
+
+Next, you need to clone the **LWM** model from its Git repository. This will download all the necessary files to your local system.
+
+```bash
+
 # Step 1: Clone the model repository (if not already cloned)
 model_repo_url = "https://huggingface.co/sadjadalikhani/lwm"
 model_repo_dir = "./LWM"
@@ -98,7 +156,7 @@ if not os.path.exists(model_repo_dir):
 
 ### 5. **Clone the Desired Datasets**
 
-LWM is designed to process datasets from various environments, such as the **DeepMIMO** dataset. Before running inference, ensure the datasets are cloned into your local environment. Below is an overview of some of the available datasets and their respective links:
+Before proceeding with tokenization and data processing, the **DeepMIMO** dataset—or any dataset generated using the operational settings outlined below—must first be loaded. The table below provides a list of available datasets and their respective links for further details:
 
 📊 **Dataset Overview**
 
@@ -106,28 +164,61 @@ LWM is designed to process datasets from various environments, such as the **Dee
 |----------------|----------------------|------------------------|------------------------------------------------------------------------------------------------------------|
 | Dataset 0      | 🌆 Denver             | 1354                   | [DeepMIMO City Scenario 18](https://www.deepmimo.net/scenarios/deepmimo-city-scenario18/)                   |
 | Dataset 1      | 🏙️ Indianapolis       | 3248                   | [DeepMIMO City Scenario 15](https://www.deepmimo.net/scenarios/deepmimo-city-scenario15/)                   |
+| Dataset 2      | 🌇 Oklahoma           | 3455                   | [DeepMIMO City Scenario 19](https://www.deepmimo.net/scenarios/deepmimo-city-scenario19/)                   |
+| Dataset 3      | 🌆 Fort Worth         | 1902                   | [DeepMIMO City Scenario 12](https://www.deepmimo.net/scenarios/deepmimo-city-scenario12/)                   |
+| Dataset 4      | 🌉 Santa Clara        | 2689                   | [DeepMIMO City Scenario 11](https://www.deepmimo.net/scenarios/deepmimo-city-scenario11/)                   |
+| Dataset 5      | 🌅 San Diego          | 2192                   | [DeepMIMO City Scenario 7](https://www.deepmimo.net/scenarios/deepmimo-city-scenario7/)                     |
+
+It is important to note that these six datasets were **not** used during the pre-training of the LWM model, and the high-quality embeddings produced are a testament to LWM’s robust generalization capabilities rather than overfitting.
+
+#### **Operational Settings**:
+- **Antennas at BS**: 32
+- **Antennas at UEs**: 1
+- **Subcarriers**: 32
+- **Paths**: 20
 
 ```python
-# Function to clone specific dataset scenario
-def clone_dataset_scenario(scenario_name, repo_url, model_repo_dir="./LWM", scenarios_dir="scenarios"):
-    # Logic for cloning datasets
-    pass
+# Step 2: Clone specific dataset scenario folder(s) inside the "scenarios" folder
+dataset_repo_url = "https://huggingface.co/datasets/sadjadalikhani/lwm"  # Base URL for dataset repo
+scenario_names = np.array(["city_18_denver",
+                           "city_15_indianapolis",
+                           "city_19_oklahoma", 
+                           "city_12_fortworth",
+                           "city_11_santaclara",
+                           "city_7_sandiego"]
+                          )
+scenario_idxs = np.array([3])
+selected_scenario_names = scenario_names[scenario_idxs]
+
+# Clone the requested scenario folders (this will clone every time)
+clone_dataset_scenarios(selected_scenario_names, dataset_repo_url, model_repo_dir)
 ```
 
 ---
 
-### 6. **Tokenize and Load the Model**
+### 6. **Change the working directory to LWM folder**
+```bash
+if os.path.exists(model_repo_dir):
+    os.chdir(model_repo_dir)
+    print(f"Changed working directory to {os.getcwd()}")
+else:
+    print(f"Directory {model_repo_dir} does not exist. Please check if the repository is cloned properly.")
+```
+
+---
 
-Once you have cloned the model and the datasets, you can preprocess the data and load the model.
+### 7. **Tokenize and Load the Model**
+```
 
-```python
+python
 from input_preprocess import tokenizer
 from lwm_model import lwm
+import torch
 
-# Tokenizing the dataset
-preprocessed_chs = tokenizer(selected_scenario_names=["city_18_denver"], manual_data=None, gen_raw=True)
+preprocessed_chs = tokenizer(selected_scenario_names=selected_scenario_names, 
+                             manual_data=None,
+                             gen_raw=True)
 
-# Load LWM model
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Loading the LWM model on {device}...")
 model = lwm.from_pretrained(device=device)
@@ -135,16 +226,11 @@ model = lwm.from_pretrained(device=device)
 
 ---
 
-### 7. **Perform Inference**
-
-After tokenizing the data and loading the model, you're ready to perform inference with LWM.
-
+### 8. **Perform Inference**
 ```python
 from inference import lwm_inference, create_raw_dataset
-
 input_types = ['cls_emb', 'channel_emb', 'raw']
-selected_input_type = 'cls_emb'  # Choose the type of input
-
+selected_input_type = input_types[0]
 if selected_input_type in ['cls_emb', 'channel_emb']:
     dataset = lwm_inference(preprocessed_chs, selected_input_type, model, device)
 else:
@@ -153,11 +239,11 @@ else:
 
 ---
 
-### LWM Architecture and Usage in Wireless Tasks
+### 9. **Explore the Interactive Demo**
 
-**LWM** employs a robust neural network architecture to extract valuable features from wireless channel data. Its deep layers, combined with advanced tokenization, make it suitable for challenging wireless communication tasks, such as beamforming, interference management, and channel state prediction.
+If you'd like to explore **LWM** interactively, check out the demo hosted on Hugging Face Spaces:  
+[**Try the Interactive Demo!**](https://huggingface.co/spaces/sadjadalikhani/LWM-Interactive-Demo)
 
-LWM is particularly effective in generalizing across diverse environments without overfitting, and its results are consistent across raw channel data and embeddings.
-
-By using **LWM**, you can enhance your wireless communication systems by utilizing data-driven models that outperform traditional approaches, offering faster and more accurate results.
+---
 
+Now you’re ready to dive into the world of **Large Wireless Model (LWM)**, process wireless communication datasets, and extract high-quality embeddings to fuel your research or application!