Code for building and training base modkels.

.gitignore

@@ -0,0 +1 @@
+*.json

{src → qarac}/corpora/BNCorpus.py

@@ -5,6 +5,7 @@ Created on Thu Aug 24 10:38:48 2023
 
 @author: peter
 """
+import os
 import numpy
 import numpy.random
 import nltk.corpus
@@ -16,12 +17,20 @@ def detokenize(sentences):
 class BNCorpus(object):
     
     def __init__(self,fileids=None,tokenizer=None,task=None):
-        self.bnc = nltk.corpus.reader.bnc.BNCCorpusReader('BNC/Texts',  fileids=r'[A-K]/\w*/\w*\.xml')
+        self.bnc = nltk.corpus.reader.bnc.BNCCorpusReader('/'.join([os.environ['HOME'],
+                                                                    'BNC',
+                                                                    'Texts']),  
+                                                          fileids=r'[A-K]/\w*/\w*\.xml')
         self.file_ids = self.bnc.fileids() if fileids is None else fileids
         self.n_docs = len(self.file_ids)
         self.rng = numpy.random.default_rng()
         self.tokenizer = tokenizer
         self.task = task
+        if self.tokenizer is not None:
+            self.mask = self.tokenizer.token_to_id('<mask>')
+            self.start = self.tokenizer.token_to_id('<start>')
+            self.end = self.tokenizer.token_to_id('<end>')
+            self.pad = numpy.array([self.tokenizer.token_to_id('<pad>')])
         
     def __len__(self):
         return self.n_docs
@@ -29,8 +38,8 @@ class BNCorpus(object):
     def split(self,p=0.8):
         n = int(p*self.n_docs)
         self.rng.shuffle(self.file_ids)
-        train = BNCorpus(self.fileids[:n],self.tokenizer,self.task)
-        test = BNCorpus(self.fileids[n:],self.tokenizer,self.task)
+        train = BNCorpus(self.file_ids[:n],self.tokenizer,self.task)
+        test = BNCorpus(self.file_ids[n:],self.tokenizer,self.task)
         return (train,test)
     
     def __iter__(self):
@@ -40,32 +49,47 @@ class BNCorpus(object):
             if self.task is None:
                 yield detokenize(doc)
             elif self.task=='encode':
-                yield self.endoder_example(doc)
+                yield self.encoder_example(doc)
             else:
                 yield self.decoder_example(doc)
                 
     def encoder_example(self,doc):
-        masked_sentences = []
-        sample_weights = []
-        for sentence in doc:
-            cp = sentence[:]
-            n = len(sentence)
-            weights = numpy.zeros(n)
-            k = self.rng.integers(n)
-            cp[k] = '[MASK] '
-            masked_sentences.append(cp)
-            weights[k] = 1
-            sample_weights.append(weights)
-        return (self.tokenizer.encode(detokenize(masked_sentences)),
-                self.tokenizer.encode(detokenize(doc)),
+        sentences = self.encode(doc)
+        masked_sentences = [sentence.copy()
+                            for sentence in sentences]
+        sample_weights = [numpy.zeros_like(sentence)
+                          for sentence in sentences]
+        masks = self.rng.integers([sentence.shape[0]
+                                   for sentence in sentences])
+        for (i,n) in enumerate(masks):
+            masked_sentences[i][n]=self.mask
+            sample_weights[i][n]=1
+        if sum((sentence.shape[0] for sentence in sentences))%2 ==1:
+            masked_sentences.append(self.pad)
+            sentences.append(self.pad)
+            sample_weights.append(numpy.zeros(1))
+        return (numpy.concatenate(masked_sentences),
+                numpy.concatenate(sentences),
                 numpy.concatenate(sample_weights))
+        
+            
+        
+    
+    def decoder_example(self,doc):
+        sentences = self.encode(doc)
+        before = [numpy.array([self.start])]+sentences
+        sentences.append(numpy.array([self.end]))
+        sample_weights = numpy.ones(sum([sentence.shape[0] 
+                                         for sentence in sentences]))
+        sample_weights[:4]=0
+        return (numpy.concatenate(before),
+                numpy.concatenate(sentences),
+                sample_weights)
+        
+        
+    def encode(self,doc):
+        return [numpy.array(self.tokenizer.encode(''.join(sentence)).ids)
+                for sentence in doc
+                if len(sentence)>0]
     
-    def decoder_sample(self,doc):
-        x = ['START'] + doc
-        y = doc + ['END']
-        sample_weights = [numpy.zeros(len(sentence)) if i==0
-                          else numpy.ones(len(sentence))
-                          for (i,sentence) in enumerate(y)]
-        return (self.tokenizer.encode(detokenize(x)),
-                self.tokenizer.encode(detokenize(y)),
-                numpy.concatenate(sample_weights))
+    

qarac/corpora/Batcher.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Aug 28 11:25:26 2023
+
+@author: peter
+"""
+import keras
+import tensorflow
+import tqdm
+
+class Batcher(keras.utils.Sequence):
+    
+    def __init__(self,source,batch_size=32):
+        self.batches = None
+        self.source=source
+        self.batch_size=batch_size
+        self.on_epoch_end()
+        
+    def __len__(self):
+        return len(self.batches)
+    
+    def __getitem__(self, index):
+        return self.batches[index]
+    
+    def on_epoch_end(self):
+        self.batches = []
+        n=0
+        X=[]
+        Y=[]
+        Z=[]
+        for (x,y,z) in tqdm.tqdm(self.source):
+            X.append(x)
+            Y.append(y)
+            Z.append(z)
+            n+=1
+            if n==self.batch_size:
+                self.batches.append((tensorflow.ragged.constant(X),
+                                     tensorflow.ragged.constant(Y),
+                                     tensorflow.ragged.constant(Z)))
+                n=0
+                X=[]
+                Y=[]
+                Z=[]
+        if n!=0:
+            self.batches.append((tensorflow.ragged.constant(X),
+                                 tensorflow.ragged.constant(Y),
+                                 tensorflow.ragged.constant(Z)))
+            
+    

qarac/corpora/__init__.py

@@ -0,0 +1 @@
+

{src → qarac}/models/layers/HierarchicalLogits.py

@@ -11,7 +11,8 @@ import tensorflow
 
 class LeafNode(keras.layers.Layer):
     def __init__(self):
-        self.bias = self.add_weight(shape=(1,),
+        super(LeafNode,self).__init__()
+        self.bias = self.add_weight(
                                     initializer='random_normal',
                                     trainable=True)
         
@@ -55,7 +56,7 @@ class HierarchicalLogits(keras.layers.Layer):
     def call(self,X,training=None):
         
         y=tensorflow.tensordot(X,self.normal,1)
-        result = self.concat([self.left(X)+y,self.right(X)]-y)
+        result = self.concat([self.left(X)+y,self.right(X)-y])
         return result
         
     

{src → qarac}/models/layers/HyenaLayer.py

@@ -9,15 +9,23 @@ Created on Tue Aug 22 09:34:14 2023
 import keras
 import keras_nlp
 import tensorflow
+import warnings
 
 def convolve(x,y):
-    xT = tensorflow.transpose(x,[0,2,1])
-    yT = tensorflow.transpose(y,[0,2,1])
-    z = tensorflow.signal.irfft(tensorflow.signal.rfft(xT)*tensorflow.signal.rfft(yT))
-    return tensorflow.transpose(z,[0,2,1])
-
-    
     
+    fx = tensorflow.vectorized_map(fft, x, warn=False)
+    fy = tensorflow.vectorized_map(fft, y, warn=False)
+    fz = fx*fy
+    return tensorflow.vectorized_map(ifft,fz,warn=False)
+
+@tensorflow.function    
+def fft(x):
+    return tensorflow.signal.rfft(tensorflow.transpose(x))
+ 
+@tensorflow.function   
+def ifft(x):
+    return tensorflow.transpose(tensorflow.signal.irfft(x))
+
 
 class HyenaLayer(keras.layers.Layer):
     """Keras implementation of Hyena layer. Unlike in the original paper,
@@ -40,29 +48,50 @@ class HyenaLayer(keras.layers.Layer):
         None.
 
         """
-        
+        super(HyenaLayer,self).__init__()
         self.stages = stages
         self.causal = causal
         self.data_projection = None
         self.filters = None
-        self.positional_encoding = keras_nlp.layers.SinePositionalEmbedding()
+    
+    def positional_encoding(self,X):
+        t = tensorflow.dtypes.saturate_cast(tensorflow.ragged.range(X.row_lengths()),
+                                            tensorflow.float32)
+        width = X.shape[-1]//2
+        f =10000 **tensorflow.expand_dims(-tensorflow.range(width,
+                                                            dtype=tensorflow.float32)/width,
+                                           axis=0)
+        phi = tensorflow.RaggedTensor.from_row_lengths(t.flat_values * f,
+                                                       X.row_lengths())
+
+        return tensorflow.concat([tensorflow.sin(phi),
+                                  tensorflow.cos(phi)],
+                                 axis=-1)
+        
         
     def build(self,input_shape):
-        self.data_projection = keras.layers.TimeDistributed(keras.layers.Dense((self.stages+1,input_shape[2]),
-                                                                               activation='linear'))
-        self.filters = keras.layers.TimeDistributed((self.stages,input_shape[2]),
-                                                    activation='linear')
+        width = input_shape[-1]
+        self.data_projection = self.add_weight(shape=(width,width,self.stages+1),
+                                               trainable=True)
+        self.filters = self.add_weight(shape=(width,width,self.stages),
+                                       trainable=True)
         
     def call(self,X,training=None):
-        x = self.data_projection(X)
-        f = self.filters(self.positional_encoding(X))
+        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())
         if self.causal:
             concat = keras.layers.Concatenate()
             x = concat(x,tensorflow.zeros_like(x))
             f = concat(f,tensorflow.zeros_like(f))
-        y = x[0]
+        y = x[:,:,:,0]
         for i in range(self.stages):
-            y = convolve(y,f[i])*x[i+1]
+            y = convolve(y,f[:,:,:,i])*x[:,:,:,i+1]
         if self.causal:
             for (i,n) in enumerate(X.row_lengths()):
                 y[i] = y[i,:n]

qarac/models/qarac_base_model.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Aug 23 09:50:14 2023
+
+@author: peter
+"""
+
+import keras
+import qarac.models.layers.HierarchicalLogits
+import qarac.models.layers.HyenaLayer
+
+def qarac_base_model(vocab_size,width,depth,decoder=True):
+    print('Building','decoder' if decoder else 'encoder','model with vocab size',
+          vocab_size,',',depth,'layers and vector width',width)
+    stack = [keras.layers.Input(shape=(None,),ragged=True),
+             keras.layers.Embedding(vocab_size,width,name='Embedding')]
+    for _ in range(depth):
+        stack.append(qarac.models.layers.HyenaLayer.HyenaLayer(causal=decoder))
+    #stack.append(keras.layers.TimeDistributed(qarac.models.layers.HierarchicalLogits.HierarchicalLogits(vocab_size)))
+    #stack.append(keras.layers.TimeDistributed(keras.layers.Softmax()))
+    stack.append(keras.layers.TimeDistributed(keras.layers.Dense(vocab_size,activation='softmax')))
+    return keras.models.Sequential(stack)

scripts.py

@@ -0,0 +1,50 @@
+
+import os
+import argparse
+import pickle
+import tokenizers
+import qarac.corpora.BNCorpus
+import qarac.corpora.Batcher
+import qarac.models.qarac_base_model
+import keras
+
+
+
+        
+def train_base_model(task,filename):
+    tokenizer = tokenizers.Tokenizer.from_pretrained('xlm-roberta-base')
+    tokenizer.add_special_tokens(['<start>','<end>','<pad>'])
+    tokenizer.save('/'.join([os.environ['HOME'],
+                            'QARAC',
+                            'models',
+                            'tokenizer.json']))
+    bnc = qarac.corpora.BNCorpus.BNCorpus(tokenizer=tokenizer,
+                                          task=task)
+    (train,test)=bnc.split(0.01)
+    train_data=qarac.corpora.Batcher.Batcher(train)
+    model = qarac.models.qarac_base_model.qarac_base_model(tokenizer.get_vocab_size(), 
+                                                           768, 
+                                                           12,
+                                                           task=='decode')
+    optimizer = keras.optimizers.Nadam(learning_rate=keras.optimizers.schedules.ExponentialDecay(1.0e-5, 100, 0.99))
+    model.compile(optimizer=optimizer,loss='sparse_categorical_crossentropy',metrics='accuracy')
+    model.fit(train_data,
+              epochs=100,
+              workers = 16)
+    test_data=qarac.corpora.Batcher.Batcher(test)
+    print(model.evaluate(test_data))
+    model.save(filename)
+    
+    
+    
+    
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(prog='QARAC',
+                                     description='Experimental NLP system, aimed at improving factual accuracy')
+    parser.add_argument('task')
+    parser.add_argument('-f','--filename')
+    parser.add_argument('-t','--training-task')
+    args = parser.parse_args()
+    if args.task == 'train_base_model':
+        train_base_model(args.training_task,args.filename)
+   

src/models/quarac_base_model.py

@@ -1,18 +0,0 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-Created on Wed Aug 23 09:50:14 2023
-
-@author: peter
-"""
-
-import keras
-import layers
-
-def quarac_base_model(vocab_size,width,depth,decoder=True):
-    stack = [keras.layers.Embedding(vocab_size,width)]
-    for _ in range(depth):
-        stack.append(layers.HyenaLayer(causal=decoder))
-    stack.append(keras.layers.Timedistributed(layers.HierarchicalLogits()))
-    stack.append(keras.layers.Timedistributed(keras.layers.Softmax()))
-    return keras.models.Sequential(stack)