apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

reshape for eval layer-to-layer

Browse source
This commit is contained in:
51616 2025-07-07 09:37:17 +00:00
parent 9613c65c20
commit 3f174b52f8
1 changed files with 15 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

25
src/ctx_to_lora/modeling/aggregator.py
View file

@ -48,11 +48,11 @@ class AggregatorConfig:
lora_r: int
per_rank_gen: bool
layer_to_layer_ctx_encoder: bool # = False
n_latent_queries: int # = 208
num_blocks: int
num_self_attn_per_block: int
shared_weights: bool
layer_to_layer_ctx_encoder: bool # = False # default value for backward compat
def get_aggregator_config(
@ -166,21 +166,26 @@ class Perceiver(nn.Module):
# )
if self.layer_to_layer:
ctx_features = rearrange(
ctx_features,
"1 num_layers seq_len feature_dim -> 1 (num_layers seq_len) feature_dim",
)
if ctx_attn_mask is not None:
ctx_attn_mask = repeat(
ctx_attn_mask,
"bs seq_len -> (num_layers bs) seq_len",
num_layers=self.num_layers,
)
ctx_position_ids = repeat(
ctx_position_ids,
"1 seq_len -> 1 (num_layers seq_len)",
num_layers=self.num_layers,
)
ctx_features = rearrange(
ctx_features,
"bs num_layers seq_len feature_dim -> (num_layers bs) seq_len feature_dim",
)
if ctx_position_ids is not None:
ctx_position_ids = repeat(
ctx_position_ids,
"1 seq_len -> 1 (num_layers seq_len)",
num_layers=self.num_layers,
)
ctx_features = rearrange(
ctx_features,
"1 num_layers seq_len feature_dim -> 1 (num_layers seq_len) feature_dim",
)
# print(f"ctx_features shape: {ctx_features.shape}")
# print(