GlobalAttentionPoolingHead layer

.ipynb_checkpoints/Model visualisation-checkpoint.ipynb

@@ -1,6 +0,0 @@
-{
- "cells": [],
- "metadata": {},
- "nbformat": 4,
- "nbformat_minor": 5
-}

qarac/models/layers/GlobalAttentionPoolingHead.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Sep  5 07:32:55 2023
+
+@author: peter
+"""
+
+import keras
+import tensorflow
+
+class GlobalAttentionPoolingHead(keras.layers.Layer):
+    
+    def __init__(self):
+        super(GlobalAttentionPoolingHead,self).__init__()
+        self.global_projection = None
+        self.local_projection = None
+        
+        
+    def build(self,input_shape):
+        width = input_shape[-1]
+        self.global_projection = self.add_weight('global projection',shape=(width,width))
+        self.local_projection = self.add_weight('local projection',shape=(width,width))
+        self.build=True
+    
+    @tensorflow.function
+    def project(self,X):
+        return tensorflow.tensordot(X,self.local_projection,axes=1)
+    
+    def attention_function(self,gp):
+        @tensorflow.function
+        def inner(lp):
+            return tensorflow.tensordot(lp,gp,axes=1)
+        return inner
+        
+    def call(self,X,training=None):
+        gp = tensorflow.linalg.l2_normalize(tensorflow.tensordot([tensorflow.reduce_sum(X,
+                                                                                       axis=1),
+                                                                  self.global_projection],
+                                                                 axes=1),
+                                            axis=1)
+        lp = tensorflow.linalg.l2_normalize(tensorflow.ragged.map_flat_values(self.project,
+                                                                              X),
+                                            axis=2)
+        attention = tensorflow.ragged.map_flat_values(self.attention_function(gp), 
+                                                      lp)
+        return tensorflow.reduce_sum(attention *X,
+                                     axis=1)

qarac/models/layers/HyenaLayer.py

@@ -11,22 +11,33 @@ import keras_nlp
 import tensorflow
 import warnings
 
+
 @tensorflow.function
 def convolve(x,y):
-    
-    fx = tensorflow.vectorized_map(fft, x, warn=False)
-    fy = tensorflow.vectorized_map(fft, y, warn=False)
+    xT = tensorflow.vectorized_map(tensorflow.transpose, x)
+    yT = tensorflow.vectorized_map(tensorflow.transpose, y)
+    fx = tensorflow.vectorized_map(tensorflow.signal.rfft, xT)
+    fy = tensorflow.vectorized_map(tensorflow.signal.rfft, yT)
     fz = fx*fy
-    return tensorflow.vectorized_map(ifft,fz,warn=False)
+    zT = tensorflow.vectorized_map(tensorflow.signal.irfft, fz)
+    return tensorflow.vectorized_map(tensorflow.transpose,zT)
 
-@tensorflow.function    
-def fft(x):
-    return tensorflow.signal.rfft(tensorflow.transpose(x))
+# @tensorflow.function    
+# def fft(x):
+#     return tensorflow.signal.rfft(tensorflow.transpose(x))
  
-@tensorflow.function   
-def ifft(x):
-    return tensorflow.transpose(tensorflow.signal.irfft(x))
+# @tensorflow.function   
+# def ifft(x):
+#     return tensorflow.transpose(tensorflow.signal.irfft(x))
 
+@tensorflow.function
+def pad(x):
+    return tensorflow.concat([x,tensorflow.zeros_like(x)],0)
+
+@tensorflow.function()
+def truncate(args):
+    (data,length)=args
+    return data[:length]
 
 class HyenaLayer(keras.layers.Layer):
     """Keras implementation of Hyena layer. Unlike in the original paper,
@@ -77,24 +88,31 @@ class HyenaLayer(keras.layers.Layer):
                                                trainable=True)
         self.filters = self.add_weight(shape=(width,width,self.stages),
                                        trainable=True)
+        self.built = True
+        
+    def conpute_output_shape(self,input_shape):
+        return input_shape
+        
+    @tensorflow.function
+    def project(self,x):
+        return tensorflow.tensordot(x,self.data_projection,axes=1)
+    
+    @tensorflow.function
+    def generate_filters(self,t):
+        return tensorflow.tensordot(t, self.filters,axes=1)
         
     def call(self,X,training=None):
-        x_flat = tensorflow.tensordot(X.flat_values,
-                                      self.data_projection,
-                                      axes=1)
-        f_flat = tensorflow.tensordot(self.positional_encoding(X).flat_values,
-                                      self.filters,
-                                      axes=1)
-        x = tensorflow.RaggedTensor.from_row_lengths(x_flat,X.row_lengths())
-        f = tensorflow.RaggedTensor.from_row_lengths(f_flat,X.row_lengths())
+        
+        x = tensorflow.ragged.map_flat_values(self.project, X)
+        f = tensorflow.ragged.map_flat_values(self.generate_filters,self.positional_encoding(X))
         if self.causal:
-            concat = keras.layers.Concatenate()
-            x = concat(x,tensorflow.zeros_like(x))
-            f = concat(f,tensorflow.zeros_like(f))
+            x = tensorflow.vectorize_map(pad,x)
+            f = tensorflow.vectorize_map(pad,f)
         y = x[:,:,:,0]
         for i in tensorflow.range(self.stages):
             y = convolve(y,f[:,:,:,i])*x[:,:,:,i+1]
         if self.causal:
-            for (i,n) in enumerate(X.row_lengths()):
-                y[i] = y[i,:n]
-        return y
+            y = tensorflow.vectorized_map(truncate,(y,X.row_lengths()))
+        return tensorflow.raw_ops.RaggedTensorToVariant(rt_nested_splits=y.row_splits,
+                                                        rt_dense_values=y.flat_values,
+                                                        batched_input=True)

scripts.py

@@ -7,7 +7,9 @@ import qarac.corpora.BNCorpus
 import qarac.corpora.Batcher
 import qarac.models.qarac_base_model
 import keras
+import tensorflow
 
+#tensorflow.debugging.disable_traceback_filtering()