Update lwm_model.py

lwm_model.py

@@ -10,11 +10,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-
 from inference import *
 from load_data import load_DeepMIMO_data
 from input_preprocess import *
-#from lwm_model import LWM, load_model
+from huggingface_hub import hf_hub_download
 
 
 ELEMENT_LENGTH = 16
@@ -111,51 +110,6 @@ class EncoderLayer(nn.Module):
         attn_outputs = self.norm(attn_outputs)
         enc_outputs = self.pos_ffn(attn_outputs)
         return enc_outputs, attn
-    
-# class LWM(torch.nn.Module):
-#     def __init__(self, element_length=16, d_model=64, max_len=129, n_layers=12):
-#         super().__init__()
-
-#         self.embedding = Embedding(element_length, d_model, max_len)
-#         self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])
-#         self.linear = nn.Linear(d_model, d_model)
-#         self.norm = LayerNormalization(d_model)
-
-#         embed_weight = self.embedding.proj.weight
-#         d_model, n_dim = embed_weight.size()
-#         self.decoder = nn.Linear(d_model, n_dim, bias=False)
-#         self.decoder.weight = nn.Parameter(embed_weight.transpose(0, 1))
-#         self.decoder_bias = nn.Parameter(torch.zeros(n_dim))
-
-#     @classmethod
-#     def from_pretrained(cls, ckpt_name='model_weights.pth', device='cuda'):
-#         # Define model
-#         model = cls().to(device)
-        
-#         # Download the model weights (from a remote or local repository)
-#         ckpt_path = f'https://huggingface.co/sadjadalikhani/LWM/resolve/main/{ckpt_name}'
-        
-#         # Load the model weights
-#         model.load_state_dict(torch.hub.load_state_dict_from_url(ckpt_path, map_location=device))
-#         print(f"Model loaded successfully from {ckpt_path} to {device}")
-        
-#         return model
-
-#     def forward(self, input_ids, masked_pos):
-#         output = self.embedding(input_ids)
-        
-#         for layer in self.layers:
-#             output, _ = layer(output)
-
-#         masked_pos = masked_pos.long()[:, :, None].expand(-1, -1, output.size(-1))
-#         h_masked = torch.gather(output, 1, masked_pos)
-#         h_masked = self.norm(F.relu(self.linear(h_masked)))
-#         logits_lm = self.decoder(h_masked) + self.decoder_bias
-
-#         return logits_lm, output
-
-from huggingface_hub import hf_hub_download
-import torch
 
 class LWM(torch.nn.Module):
     def __init__(self, element_length=16, d_model=64, max_len=129, n_layers=12):