From a4b741954f418eff7bb84a4c05a9b8d237c74fae Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Tue, 7 Jan 2025 11:41:28 +0000
Subject: [PATCH] add num_latent_factor to config

---
 hyperlora/configs.py        |  4 ++++
 hyperlora/modeling_utils.py | 12 ++++++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/hyperlora/configs.py b/hyperlora/configs.py
index 0d80531..c137a38 100644
--- a/hyperlora/configs.py
+++ b/hyperlora/configs.py
@@ -304,6 +304,10 @@ class AggregatorArguments:
         default=0.0,
         metadata={"help": "Attention dropout probability for Perceiver."},
     )
+    num_latent_factor: int = field(
+        default=8,
+        metadata={"help": "Number of latent factors for Perceiver."},
+    )
     num_blocks: int = field(
         default=8,
         metadata={"help": "Number of blocks for Perceiver."},
diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py
index 9575201..30a1e69 100644
--- a/hyperlora/modeling_utils.py
+++ b/hyperlora/modeling_utils.py
@@ -58,6 +58,7 @@ class AggregatorConfig:
     num_self_attends_per_block: int = 16
     self_attention_widening_factor: int = 4
     cross_attention_widening_factor: int = 1
+    num_latent_factor: int = 8
 
 
 def get_aggregator_config(
@@ -111,14 +112,21 @@ class Perceiver(nn.Module):
     """perceiver w/ bottleneck size = n_modules * n_layers"""
 
     def __init__(
-        self, feature_size, output_size, num_layers, num_modules, *args, **kwargs
+        self,
+        feature_size,
+        output_size,
+        num_layers,
+        num_modules,
+        num_latent_factor,
+        *args,
+        **kwargs,
     ):
         super().__init__()
         self.num_layers = num_layers
         self.num_modules = num_modules
         self.config = PerceiverConfig(
             d_model=feature_size,  # + num_bands
-            num_latents=num_layers * num_modules * 8,
+            num_latents=num_layers * num_modules * num_latent_factor,
             d_latents=output_size,
             # attention_probs_dropout_prob=0.0,
             # num_blocks=8,