Using torch.nn.CosineSimilarity to simplify code

qarac/models/QaracTrainerModel.py

@@ -10,7 +10,7 @@ import torch
 import qarac.models.QaracEncoderModel
 import qarac.models.QaracDecoderModel
 
-EPSILON=torch.tensor(1.0e-12)
+
 
 class QaracTrainerModel(torch.nn.Module):
     
@@ -39,6 +39,7 @@ class QaracTrainerModel(torch.nn.Module):
         self.decoder = qarac.models.QaracDecoderModel.QaracDecoderModel(base_model_path,
                                                                         config,
                                                                         tokenizer)
+        self.cosine = torch.nn.CosineSimilarity(dim=2,eps=1.0e-12)
         
     def forward(self,
                 all_text,
@@ -92,15 +93,9 @@ class QaracTrainerModel(torch.nn.Module):
         reasoning = self.decoder((self.answer_encoder(proposition0)
                                              +self.answer_encoder(proposition1),
                                              conclusion_offset))
-        s0vec = self.answer_encoder(statement0)
-        s0norm = torch.maximum(torch.linalg.vector_norm(s0vec,
-                                                        dim=1,
-                                                        keepdim=True),EPSILON)
-        s0 = s0vec/s0norm
-        s1vec = self.answer_encoder(statement1)
-        s1norm = torch.maximum(torch.linalg.vector_norm(s1vec,
-                                                        dim=1,
-                                                        keepdim=True),EPSILON)
-        s1 = s1vec/s1norm
-        consistency = torch.einsum('ij,ij->i',s0,s1)
+        s0 = self.answer_encoder(statement0)
+        
+        s1 = self.answer_encoder(statement1)
+        
+        consistency = self.cosine(s0,s1)
         return (encode_decode,question_answering,reasoning,consistency)

qarac/models/layers/GlobalAttentionPoolingHead.py

@@ -8,7 +8,6 @@ Created on Tue Sep  5 07:32:55 2023
 
 import torch
 
-EPSILON = torch.tensor(1.0e-12)   
 
 class GlobalAttentionPoolingHead(torch.nn.Module):
     
@@ -29,6 +28,7 @@ class GlobalAttentionPoolingHead(torch.nn.Module):
         super(GlobalAttentionPoolingHead,self).__init__()
         self.global_projection = torch.nn.Linear(size,size,bias=False)
         self.local_projection = torch.nn.Linear(size,size,bias=False)
+        self.cosine = torch.nn.CosineSimilarity(dim=2,eps=1.0e-12)
         
     
         
@@ -55,16 +55,9 @@ class GlobalAttentionPoolingHead(torch.nn.Module):
         else:
             attention_mask = attention_mask.unsqueeze(2)
         Xa = X*attention_mask
-        sigma = torch.sum(Xa,dim=1)
-        psigma = self.global_projection(sigma)
-        nsigma = torch.maximum(torch.linalg.vector_norm(psigma,
-                                                        dim=1,
-                                                        keepdim=True),EPSILON)
-        gp = psigma/nsigma
-        loc = self.local_projection(Xa)
-        nloc = torch.maximum(torch.linalg.vector_norm(loc,
-                                                      dim=2,
-                                                      keepdim=True),EPSILON)
-        lp = loc/nloc
-        attention = torch.einsum('ijk,ik->ij',lp,gp)
+        sigma = torch.sum(Xa,dim=1,keepdim=True)
+        gp = self.global_projection(sigma)
+        lp = self.local_projection(Xa)
+        
+        attention = self.cosine(lp,gp)
         return torch.einsum('ij,ijk->ik',attention,Xa)

scripts.py

@@ -149,6 +149,7 @@ def train_models(path):
                                  X['conclusion_offset'],
                                  X['statement0'],
                                  X['statement1'])
+            print([y.shape for y in prediction])
             loss = loss_fn(prediction,Y)
             loss.backward()
             optimizer.step()
@@ -411,16 +412,9 @@ def test_consistency(path):
     s1_attn = torch.not_equal(s1_in,
                               pad_token)
     s0_vec = encoder(s0_in,attention_mask=s0_attn)
-    s0_norm = torch.maximum(torch.linalg.vector_norm(s0_vec,
-                                                     dim=1,
-                                                     keepdim=True),EPSILON)
-    s0 = s0_vec/s0_norm
     s1_vec = encoder(s1_in,attention_mask=s1_attn)
-    s1_norm = torch.maximum(torch.linalg.vector_norm(s1_vec,
-                                                     dim=1,
-                                                     keepdim=True),EPSILON)
-    s1 = s1_vec/s1_norm
-    consistency = torch.einsum('ij,ij->i',s0,s1).numpy()
+    cosine = torch.nn.CosineSimilarity(dim=2,eps=1.0e-12)
+    consistency = cosine(s0_vec,s1_vec).numpy()
     results = pandas.DataFrame({'label':data['gold_label'],
                                 'score':consistency})
     third = 1.0/3.0