.

Dockerfile

@@ -4,7 +4,7 @@ FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
 
 WORKDIR /app
 
-COPY . /app
+
 
 SHELL ["/bin/bash", "-c"]
 
@@ -93,6 +93,8 @@ RUN ls /usr/local/lib/python3.10/dist-packages/lgatr/layers
 RUN cd /opt/pepr && git clone https://github.com/cms-pepr/pytorch_cmspepr
 RUN cd /opt/pepr/pytorch_cmspepr/ && python3 -m pip install .
 
+COPY . /app
+
 # entrypoint run app.py with python
 EXPOSE 7860
 ENV GRADIO_SERVER_NAME="0.0.0.0"

src/1models/LGATr/lgatr.py

@@ -1,101 +1,5 @@
-import torch
-try:
-    from lgatr import LGATr, SelfAttentionConfig, MLPConfig
-    from lgatr.interface import embed_vector, extract_scalar, extract_vector
-    #from lgatr.interface.spurions import get_spurions
-
-    def embed_spurions(
-            beam_reference,
-            add_time_reference,
-            two_beams=True,
-            add_xzplane=False,
-            add_yzplane=False,
-            device="cpu",
-            dtype=torch.float32,
-    ):
-        """
-        Construct a list of reference multivectors/spurions for symmetry breaking
-
-        Parameters
-        ----------
-        beam_reference: str
-            Different options for adding a beam_reference
-            Options: "lightlike", "spacelike", "timelike", "xyplane"
-        add_time_reference: bool
-            Whether to add the time direction as a reference to the network
-        two_beams: bool
-            Whether we only want (x, 0, 0, 1) or both (x, 0, 0, +/- 1) for the beam
-        add_xzplane: bool
-            Whether to add the x-z-plane as a reference to the network
-        add_yzplane: bool
-            Whether to add the y-z-plane as a reference to the network
-        device
-        dtype
-
-        Returns
-        -------
-        spurions: torch.tensor with shape (n_spurions, 16)
-            spurion embedded as multivector object
-        """
-        kwargs = {"device": device, "dtype": dtype}
-
-        if beam_reference in ["lightlike", "spacelike", "timelike"]:
-            # add another 4-momentum
-            if beam_reference == "lightlike":
-                beam = [1, 0, 0, 1]
-            elif beam_reference == "timelike":
-                beam = [2 ** 0.5, 0, 0, 1]
-            elif beam_reference == "spacelike":
-                beam = [0, 0, 0, 1]
-            beam = torch.tensor(beam, **kwargs).reshape(1, 4)
-            beam = embed_vector(beam)
-            if two_beams:
-                beam2 = beam.clone()
-                beam2[..., 4] = -1  # flip pz
-                beam = torch.cat((beam, beam2), dim=0)
-
-        elif beam_reference == "xyplane":
-            # add the x-y-plane, embedded as a bivector
-            # convention for bivector components: [tx, ty, tz, xy, xz, yz]
-            beam = torch.zeros(1, 16, **kwargs)
-            beam[..., 8] = 1
-
-        elif beam_reference is None:
-            beam = torch.empty(0, 16, **kwargs)
-
-        else:
-            raise ValueError(f"beam_reference {beam_reference} not implemented")
-
-        if add_xzplane:
-            # add the x-z-plane, embedded as a bivector
-            xzplane = torch.zeros(1, 16, **kwargs)
-            xzplane[..., 10] = 1
-        else:
-            xzplane = torch.empty(0, 16, **kwargs)
-
-        if add_yzplane:
-            # add the y-z-plane, embedded as a bivector
-            yzplane = torch.zeros(1, 16, **kwargs)
-            yzplane[..., 9] = 1
-        else:
-            yzplane = torch.empty(0, 16, **kwargs)
-
-        if add_time_reference:
-            time = [1, 0, 0, 0]
-            time = torch.tensor(time, **kwargs).reshape(1, 4)
-            time = embed_vector(time)
-        else:
-            time = torch.empty(0, 16, **kwargs)
-
-        spurions = torch.cat((beam, xzplane, yzplane, time), dim=-2)
-        return spurions
-except:
-    print("Failed importing, trying importing GATr")
-    from lgatr import SelfAttentionConfig, MLPConfig
-    from lgatr import GATr as LGATr
-    from lgatr.interface import embed_vector, extract_scalar, extract_vector, embed_spurions
-
-
+from lgatr import GATr, SelfAttentionConfig, MLPConfig
+from lgatr.interface import embed_vector, extract_scalar, embed_spurions, extract_vector
 import torch
 import torch.nn as nn
 from xformers.ops.fmha import BlockDiagonalMask
@@ -115,7 +19,7 @@ class LGATrModel(torch.nn.Module):
         self.n_scalars_out = n_scalars_out
         self.obj_score = obj_score
         self.global_features_copy = global_featuers_copy
-        self.gatr = LGATr(
+        self.gatr = GATr(
             in_mv_channels=3,
             out_mv_channels=1,
             hidden_mv_channels=hidden_mv_channels,