Update nets/model.py

nets/model.py

@@ -4,6 +4,7 @@ from nets.decoder import Decoder
 from nets.projections import Projections
 from nets.encoder import Encoder
 
+
 class Model(nn.Module):
     def __init__(self, input_size, embedding_size,
                  num_heads=8, num_layers=4, ff_hidden=250, *args, **kwargs):
@@ -11,34 +12,30 @@ class Model(nn.Module):
 
         self.embedding_size = embedding_size
 
-        # ---------------------- Encoder ------------------------
-        # Takes input features like: [lng, lat, start_time, end_time, depot_flag]
+        # ----------- Encoder -----------
         self.encoder = Encoder(
             n_heads=num_heads,
             embed_dim=embedding_size,
             n_layers=num_layers,
             feed_forward_hidden=ff_hidden,
-            node_dim=input_size  # ← هذا هو عدد الخصائص للعقدة (5 غالبًا)
+            node_dim=input_size
         )
 
-        # ---------------------- Decoder ------------------------
-        # سابقًا كانت ثابتة 4*embedding + 1, هنا نخليها أكثر مرونة
-        decoder_input_dim = 3 * embedding_size + 2  # [encoder_output, mask, fleet_info] تقريبًا
-
+        # ----------- Decoder -----------
+        decoder_input_dim = 3 * embedding_size + 2  # example: [encoder output] + [mask info] + [fleet info]
         self.decoder = Decoder(
             input_size=decoder_input_dim,
             embedding_size=embedding_size,
             num_heads=num_heads
         )
 
-        # ---------------------- Attention Projections ------------------------
+        # ----------- Attention Projections -----------
         self.projections = Projections(
             n_heads=num_heads,
             embed_dim=embedding_size
         )
 
-        # ---------------------- Fleet Attention ------------------------
-        # هذه تستخدم بيانات الأسطول مثل [start_time, location_embedding]
+        # ----------- Fleet Attention Encoder (Optional) -----------
         self.fleet_attention = Encoder(
             n_heads=num_heads,
             embed_dim=embedding_size,