patchfm 2.2.2__tar.gz → 2.2.5__tar.gz

{patchfm-2.2.2/src/patchfm.egg-info → patchfm-2.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchfm
-Version: 2.2.2
+Version: 2.2.5
 Summary: a Foundation Model for Univariate Time Series Forecasting
 Author-email: Samy-Melwan Vilhes <samy-melwan.vilhes@insa-rouen.fr>
 License: MIT License

{patchfm-2.2.2 → patchfm-2.2.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "patchfm"
-version = "2.2.2"
+version = "2.2.5"
 authors = [
   { name="Samy-Melwan Vilhes", email="samy-melwan.vilhes@insa-rouen.fr" },
 ]

{patchfm-2.2.2 → patchfm-2.2.5}/src/patchfm/configs/model_config.py

@@ -8,6 +8,7 @@ class PatchFMConfig:
     d_model: int = 2048
     n_heads: int = 32
     n_layers_encoder: int = 6
+    use_xsa: bool = True
     quantiles: list[float] = field(default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
     
     compile: bool = True

{patchfm-2.2.2 → patchfm-2.2.5}/src/patchfm/inference/forecaster.py

@@ -18,12 +18,18 @@ class Forecaster(nn.Module):
         self.d_model = config["d_model"]
         self.n_heads = config["n_heads"]
         self.n_layers_encoder = config["n_layers_encoder"]
+        self.use_xsa = config["use_xsa"]
         self.quantiles = config["quantiles"]
         self.n_quantiles = len(self.quantiles)
         self.max_patches = self.max_seq_len // self.patch_len
 
-        print("Loading base model from HuggingFace Hub...")
-        base_model = PatchFM.from_pretrained("vilhess/PatchFM")
+        if not self.use_xsa:
+            print("Loading base model from HuggingFace Hub...")
+            base_model = PatchFM.from_pretrained("vilhess/PatchFM")
+        else:
+            print("Loading XSA model from HuggingFace Hub...")
+            base_model = PatchFM.from_pretrained("vilhess/PatchFM-XSA")
+
         self._init_from_base(base_model)
 
         self.eval()
@@ -46,7 +52,8 @@ class Forecaster(nn.Module):
         self.transformer_encoder = TransformerEncoder(
             d_model=self.d_model, 
             n_heads=self.n_heads, 
-            n_layers=self.n_layers_encoder
+            n_layers=self.n_layers_encoder,
+            use_xsa=self.use_xsa
         )
         self.proj_output = ResidualBlock(
             in_dim=self.d_model, 
@@ -69,15 +76,7 @@ class Forecaster(nn.Module):
         # Default horizon = patch_len
         forecast_horizon = forecast_horizon or self.patch_len
 
-        x = self.converter.convert(x)
-        assert x.ndim in (1, 2), f"Input dimension must be 1D (time) or 2D (batch, time), got {x.ndim}D."
-
-        batch_dim=True
-        if x.ndim != 2:
-            x = x.unsqueeze(0)
-            batch_dim=False
         bs, ws = x.size()
-
         x = x.to(self.device)
 
         if ws > self.max_seq_len:
@@ -153,26 +152,39 @@ class Forecaster(nn.Module):
             pred_median = torch.zeros_like(pred_median)
             pred_quantiles = torch.zeros_like(pred_quantiles)
         
-        if not batch_dim:
-            pred_median = pred_median.squeeze(0)
-            pred_quantiles = pred_quantiles.squeeze(0)   
-
-        pred_median, pred_quantiles = self.converter.deconvert(pred_median, pred_quantiles)
         return pred_median, pred_quantiles
 
     def __call__(self, context: torch.Tensor, forecast_horizon: int | None = None, quantiles: list[float] | None = None, flip_equivariance: bool = False) -> torch.Tensor:
+
+        x = self.converter.convert(context)
+        assert x.ndim in (1, 2), f"Input dimension must be 1D (time) or 2D (batch, time), got {x.ndim}D."
+
+        batch_dim=True
+        if x.ndim != 2:
+            x = x.unsqueeze(0)
+            batch_dim=False
+
+        bs, ws = x.size()
+
         if flip_equivariance:
             print("Flip equivariance enabled: forecast = (f(x) - f(-x)) / 2. This requires multiplying by 2 the batch size (Reverso: Efficient Time Series Foundation Models for Zero-shot Forecasting).")
             bs = context.size(0)
-            context_flipped = -context
-            concat_context = torch.cat([context, context_flipped], dim=0)
-            pred_median_full, pred_quantiles_full = self.auto_regressive_quantile_decoding(concat_context, forecast_horizon, quantiles)
+            x_flipped = -x
+            concat_x = torch.cat([x, x_flipped], dim=0)
+            pred_median_full, pred_quantiles_full = self.auto_regressive_quantile_decoding(concat_x, forecast_horizon, quantiles)
             pred_median, pred_quantiles = pred_median_full[:bs], pred_quantiles_full[:bs]
             pred_median2, pred_quantiles2 = pred_median_full[bs:], pred_quantiles_full[bs:]
             pred_median = (pred_median - pred_median2) / 2
             pred_quantiles = (pred_quantiles - flip_last_dim(pred_quantiles2)) / 2
         else:
-            pred_median, pred_quantiles = self.auto_regressive_quantile_decoding(context, forecast_horizon, quantiles)
+            pred_median, pred_quantiles = self.auto_regressive_quantile_decoding(x, forecast_horizon, quantiles)
+
+        if not batch_dim:
+            pred_median = pred_median.squeeze(0)
+            pred_quantiles = pred_quantiles.squeeze(0)   
+
+        pred_median, pred_quantiles = self.converter.deconvert(pred_median, pred_quantiles)
+
         return pred_median, pred_quantiles
     
     def clear_cache(self):

{patchfm-2.2.2 → patchfm-2.2.5}/src/patchfm/inference/modules.py

@@ -185,7 +185,7 @@ class ResidualBlock(nn.Module):
         return out
 
 class MultiHeadAttention(nn.Module):
-    def __init__(self, d_model, n_heads, last=False):
+    def __init__(self, d_model, n_heads, last=False, use_xsa=False):
         super().__init__()
         assert d_model%n_heads==0, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
 
@@ -204,7 +204,8 @@ class MultiHeadAttention(nn.Module):
         self.v_cache = None
 
         self.last = last
-    
+        self.use_xsa = use_xsa
+
     def forward(self, q):
         bs, context, dim = q.size()
         offset = 0
@@ -239,6 +240,13 @@ class MultiHeadAttention(nn.Module):
 
         values = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
 
+        if self.use_xsa:
+            if v.size(2) != values.size(2):
+                v = v[:, :, -values.size(2) :, :]
+                
+            vn = torch.nn.functional.normalize(v, dim=-1)
+            values = values - (values * vn).sum(dim=-1, keepdim=True) * vn
+
         values = values.transpose(1, 2).reshape(bs, -1, dim)
         values = self.out_proj(values)
         return values
@@ -267,10 +275,10 @@ class FeedForward(nn.Module):
 
 
 class TransformerEncoderLayer(nn.Module):
-    def __init__(self, d_model, n_heads, last=False):
+    def __init__(self, d_model, n_heads, last=False, use_xsa=False):
         super().__init__()
         self.ln1 = nn.LayerNorm(d_model)
-        self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, last=last)
+        self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, last=last, use_xsa=use_xsa)
         self.ln2 = nn.LayerNorm(d_model)
         self.ff = FeedForward(d_model=d_model)
     
@@ -281,15 +289,15 @@ class TransformerEncoderLayer(nn.Module):
         return out
     
 class TransformerEncoder(nn.Module):
-    def __init__(self, d_model, n_heads, n_layers):
+    def __init__(self, d_model, n_heads, n_layers, use_xsa=False):
         super().__init__()
         self.layers = nn.ModuleList(
             [
-                TransformerEncoderLayer(d_model=d_model, n_heads=n_heads)
+                TransformerEncoderLayer(d_model=d_model, n_heads=n_heads, use_xsa=use_xsa)
                 for _ in range(n_layers-1)
             ]
         )
-        self.layers.append(TransformerEncoderLayer(d_model=d_model, n_heads=n_heads, last=True))
+        self.layers.append(TransformerEncoderLayer(d_model=d_model, n_heads=n_heads, last=True, use_xsa=use_xsa))
         self.norm = nn.LayerNorm(d_model)
 
     def forward(self, x):
@@ -306,6 +314,7 @@ class PatchFM(nn.Module, PyTorchModelHubMixin):
         self.d_model = config["d_model"]
         self.n_heads = config["n_heads"]
         self.n_layers_encoder = config["n_layers_encoder"]
+        self.use_xsa = config["use_xsa"]
         self.quantiles = config["quantiles"]
         self.n_quantiles = len(self.quantiles)
 
@@ -319,7 +328,8 @@ class PatchFM(nn.Module, PyTorchModelHubMixin):
         self.transformer_encoder = TransformerEncoder(
             d_model=self.d_model, 
             n_heads=self.n_heads, 
-            n_layers=self.n_layers_encoder
+            n_layers=self.n_layers_encoder, 
+            use_xsa=self.use_xsa
         )
         self.proj_output = ResidualBlock(
             in_dim=self.d_model, 

{patchfm-2.2.2 → patchfm-2.2.5/src/patchfm.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchfm
-Version: 2.2.2
+Version: 2.2.5
 Summary: a Foundation Model for Univariate Time Series Forecasting
 Author-email: Samy-Melwan Vilhes <samy-melwan.vilhes@insa-rouen.fr>
 License: MIT License