add attention mechanisms

Author: thierrymoudiki

examples/attention.py

@@ -0,0 +1,103 @@
+import os 
+import nnetsauce as ns 
+import numpy as np 
+import jax.numpy as jnp
+from nnetsauce.attention import AttentionMechanism
+from sklearn.datasets import load_diabetes, fetch_california_housing
+from sklearn.model_selection import train_test_split
+from sklearn.linear_model import Ridge
+from sklearn.ensemble import ExtraTreesRegressor, RandomForestRegressor
+from time import time 
+
+print(f"\n ----- Running: {os.path.basename(__file__)}... ----- \n")
+
+# Set random seed for reproducibility
+np.random.seed(42)
+
+# Example 1: Univariate time series with temporal attention
+print("=" * 50)
+print("Example 1: Univariate Time Series")
+print("=" * 50)
+batch_size, seq_len, input_dim = 32, 10, 1
+x_univariate = jnp.array(np.random.randn(batch_size, seq_len, input_dim))
+
+attention = AttentionMechanism(input_dim=input_dim, hidden_dim=32, num_heads=4)
+context, weights = attention(x_univariate, attention_type='temporal')
+
+print(f"Input shape: {x_univariate.shape}")
+print(f"Context shape: {context.shape}")
+print(f"Attention weights shape: {weights.shape}")
+print(f"Sample attention weights (first batch): {np.array(weights[0])}")
+
+# Example 2: Tabular data with feature attention
+print("\n" + "=" * 50)
+print("Example 2: Tabular Data with Feature Attention")
+print("=" * 50)
+batch_size, num_features = 32, 10
+x_tabular = jnp.array(np.random.randn(batch_size, num_features))
+
+attention_tab = AttentionMechanism(input_dim=num_features, hidden_dim=32)
+output, feature_weights = attention_tab(x_tabular, attention_type='feature')
+
+print(f"Input shape: {x_tabular.shape}")
+print(f"Output shape: {output.shape}")
+print(f"Feature weights shape: {feature_weights.shape}")
+print(f"Feature importance (first batch): {np.array(feature_weights[0])}")
+
+# Example 3: Multi-head attention on sequences
+print("\n" + "=" * 50)
+print("Example 3: Multi-Head Attention")
+print("=" * 50)
+batch_size, seq_len, input_dim = 16, 8, 16
+x_seq = jnp.array(np.random.randn(batch_size, seq_len, input_dim))
+
+attention_mha = AttentionMechanism(input_dim=input_dim, hidden_dim=64, num_heads=8)
+output_mha, weights_mha = attention_mha(x_seq, attention_type='multi_head')
+
+print(f"Input shape: {x_seq.shape}")
+print(f"Output shape: {output_mha.shape}")
+print(f"Attention weights shape (with heads): {weights_mha.shape}")
+
+# Example 4: Cross-attention
+print("\n" + "=" * 50)
+print("Example 4: Cross-Attention")
+print("=" * 50)
+batch_size = 16
+query_seq = jnp.array(np.random.randn(batch_size, 5, input_dim))
+kv_seq = jnp.array(np.random.randn(batch_size, 10, input_dim))
+
+cross_output, cross_weights = attention_mha(
+    None, 
+    attention_type='cross',
+    query=query_seq,
+    key_value=kv_seq
+)
+
+print(f"Query shape: {query_seq.shape}")
+print(f"Key-Value shape: {kv_seq.shape}")
+print(f"Cross-attention output shape: {cross_output.shape}")
+print(f"Cross-attention weights shape: {cross_weights.shape}")
+
+# Example 5: Context Vector Attention
+print("\n" + "=" * 50)
+print("Example 5: Context Vector Attention")
+print("=" * 50)
+batch_size, seq_len, input_dim = 32, 15, 8
+x_context = jnp.array(np.random.randn(batch_size, seq_len, input_dim))
+
+attention_ctx = AttentionMechanism(input_dim=input_dim, hidden_dim=64)
+context_output, context_weights = attention_ctx(x_context, attention_type='context_vector')
+
+print(f"Input shape: {x_context.shape}")
+print(f"Context output shape: {context_output.shape}")
+print(f"Context attention weights shape: {context_weights.shape}")
+print(f"Sample context weights (first batch): {np.array(context_weights[0])}")
+print(f"\nNote: Context vector attention produces a fixed-size global representation")
+print(f"regardless of input sequence length, making it ideal for classification tasks.")
+
+# Demonstrate JAX's JIT compilation benefit
+print("\n" + "=" * 50)
+print("JAX Performance Benefits")
+print("=" * 50)
+print("All methods are JIT-compiled for fast execution!")
+print("JAX provides automatic differentiation and GPU/TPU acceleration.")

examples/conformal_simulation2.py

@@ -0,0 +1,96 @@
+import os 
+import nnetsauce as ns 
+import matplotlib.pyplot as plt 
+import numpy as np 
+import warnings
+from sklearn.datasets import fetch_california_housing, load_diabetes
+from sklearn.model_selection import train_test_split
+from sklearn.linear_model import BayesianRidge, ARDRegression, RidgeCV
+from sklearn.ensemble import ExtraTreesRegressor
+from time import time 
+
+
+# # 2 - Useful plotting functions
+
+
+
+warnings.filterwarnings('ignore')
+
+split_color = 'green'
+split_color2 = 'tomato'
+local_color = 'gray'
+
+def plot_func(x,
+              y,
+              y_u=None,
+              y_l=None,
+              pred=None,
+              shade_color="",
+              method_name="",
+              title=""):
+
+    fig = plt.figure()
+
+    plt.plot(x, y, 'k.', alpha=.3, markersize=10,
+             fillstyle='full', label=u'Test set observations')
+
+    if (y_u is not None) and (y_l is not None):
+        plt.fill(np.concatenate([x, x[::-1]]),
+                 np.concatenate([y_u, y_l[::-1]]),
+                 alpha=.3, fc=shade_color, ec='None',
+                 label = method_name + ' Prediction interval')
+
+    if pred is not None:
+        plt.plot(x, pred, 'k--', lw=2, alpha=0.9,
+                 label=u'Predicted value')
+
+    #plt.ylim([-2.5, 7])
+    plt.xlabel('$X$')
+    plt.ylabel('$Y$')
+    plt.legend(loc='upper right')
+    plt.title(title)
+
+    plt.show()
+
+
+# # 3 - Examples of use
+
+
+
+data = fetch_california_housing()
+X = data.data
+y= data.target
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = .2, random_state = 123)
+
+
+# ## RidgeCV 
+"""
+- 'bootstrap': Bootstrap resampling.
+- 'kde': Kernel Density Estimation.
+- 'ecdf': Empirical CDF-based sampling.
+- 'permutation': Permutation resampling.
+- 'smooth_bootstrap': Smoothed bootstrap with added noise.
+"""                      
+
+for type_pi in ('bootstrap', 'kde', 'ecdf', 'permutation', 'smooth_bootstrap'):
+    print(f"\n\n### type_pi = {type_pi} ###\n")
+    regr1 = ns.PredictionInterval(RidgeCV(),
+                            replications=100,
+                            type_pi=type_pi) # 5 hidden nodes, ReLU activation function
+    regr1.fit(X_train, y_train)
+    start = time()
+    preds1 = regr1.predict(X_test, return_pi=True)
+    print(f"Elapsed: {time() - start}s")
+    print(f"coverage_rate conformalized QRNN RidgeCV: {np.mean((preds1[2]<=y_test)*(preds1[3]>=y_test))}")
+    print(f"predictive simulations: {preds1[1]}")
+
+    max_idx = 50
+    plot_func(x = range(max_idx),
+            y = y_test[0:max_idx],
+            y_u = preds1.upper[0:max_idx],
+            y_l = preds1.lower[0:max_idx],
+            pred = preds1.mean[0:max_idx],
+            shade_color=split_color2,
+            title = f"conformalized QRNN RidgeCV ({max_idx} first points in test set)")
+
+

nnetsauce/__init__.py

@@ -1,3 +1,4 @@
+from .attention import AttentionMechanism
 from .base.base import Base
 from .base.baseRegressor import BaseRegressor
 from .boosting.adaBoostClassifier import AdaBoostClassifier
@@ -42,6 +43,7 @@
 
 __all__ = [
     "AdaBoostClassifier",
+    "AttentionMechanism",
     "Base",
     "BaseRegressor",
     "BayesianRVFLRegressor",

nnetsauce/attention/__init__.py

@@ -0,0 +1,3 @@
+from .attention import AttentionMechanism
+
+__all__ = ["AttentionMechanism"]