From 10ec27c85b27af291b4c4e92e9bfac241900ad9c Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Thu, 13 Feb 2025 21:20:36 +0530 Subject: [PATCH 01/10] Update encoder.py --- gliner/modeling/encoder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gliner/modeling/encoder.py b/gliner/modeling/encoder.py index b5ac09a0..ed54a2ba 100644 --- a/gliner/modeling/encoder.py +++ b/gliner/modeling/encoder.py @@ -97,7 +97,7 @@ def forward(self, *args, **kwargs): output_hidden_states = True else: output_hidden_states = False - output = self.model(*args, output_hidden_states = output_hidden_states, + output = self.model(*args, #output_hidden_states = output_hidden_states, return_dict = True, **kwargs) if self.config.fuse_layers: encoder_layer = self.layers_fuser(output.hidden_states) @@ -165,4 +165,4 @@ def forward(self, input_ids, attention_mask, token_embeddings = self.encode_text(input_ids, attention_mask, *args, **kwargs) labels_embeddings = self.encode_labels(labels_input_ids, labels_attention_mask, *args, **kwargs) - return token_embeddings, labels_embeddings \ No newline at end of file + return token_embeddings, labels_embeddings From 99f8c9a1b372833170eecda316bb129402e7820c Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Thu, 13 Feb 2025 21:40:23 +0530 Subject: [PATCH 02/10] Update custom_train.py --- custom_train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/custom_train.py b/custom_train.py index 99df88ec..693841b4 100644 --- a/custom_train.py +++ b/custom_train.py @@ -183,7 +183,7 @@ def setup_model_and_optimizer(self, rank=None, device=None): model.module.resize_token_embeddings([self.model_config.ent_token, self.model_config.sep_token], set_class_token_index = False, add_tokens_to_tokenizer=False) - optimizer = self.create_optimizer(model.model) + optimizer = self.create_optimizer(model.module if isinstance(model, DDP) else model) if self.compile_model: model.compile_for_training() @@ -352,4 +352,4 @@ def create_parser(): trainer = Trainer(config, allow_distributed=args.allow_distributed, compile_model = args.compile_model, device='cuda' if torch.cuda.is_available() else 'cpu') - trainer.run() \ No newline at end of file + trainer.run() From 2d973ce4f2ef00f4d0a534687fa320028b7172fe Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Thu, 13 Feb 2025 21:44:11 +0530 Subject: [PATCH 03/10] Update custom_train.py --- custom_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/custom_train.py b/custom_train.py index 693841b4..152cac16 100644 --- a/custom_train.py +++ b/custom_train.py @@ -208,7 +208,7 @@ def train_dist(self, rank, world_size, dataset): sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False) - train_loader = self.create_dataloader(dataset, model.data_processor, sampler=sampler, shuffle=False) + train_loader = train_loader = self.create_dataloader(dataset, model.module.data_processor if isinstance(model, DDP) else model.data_processor, sampler=sampler, shuffle=False) num_steps = self.config.num_steps // world_size From a81645df0c7e0484cbdb82fe0a54876a15ea3495 Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Thu, 13 Feb 2025 21:52:03 +0530 Subject: [PATCH 04/10] Update custom_train.py --- custom_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/custom_train.py b/custom_train.py index 152cac16..33440f72 100644 --- a/custom_train.py +++ b/custom_train.py @@ -194,7 +194,7 @@ def create_dataloader(self, dataset, data_processor, sampler=None, shuffle=True) # dataset = GLiNERDataset(dataset, config = self.config, data_processor=self.data_processor) # collator = DataCollatorWithPadding(self.config) collator = DataCollator(self.config, data_processor=data_processor, prepare_labels=True) - data_loader = DataLoader(dataset, batch_size=self.config.train_batch_size, num_workers=12, + data_loader = DataLoader(dataset, batch_size=self.config.train_batch_size, num_workers=2, shuffle=shuffle, collate_fn=collator, sampler=sampler) return data_loader From 4c9d474cef6b02a977a08a4d9d0c0eea75c55c08 Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Thu, 13 Feb 2025 22:04:34 +0530 Subject: [PATCH 05/10] Update custom_train.py --- custom_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/custom_train.py b/custom_train.py index 33440f72..2a8efe98 100644 --- a/custom_train.py +++ b/custom_train.py @@ -178,7 +178,7 @@ def setup_model_and_optimizer(self, rank=None, device=None): set_class_token_index = False, add_tokens_to_tokenizer=False) if rank is not None: - model = DDP(model, device_ids=[rank], output_device=rank, find_unused_parameters=False) + model = DDP(model, device_ids=[rank], output_device=rank, find_unused_parameters=True) if self.config.labels_encoder is None: model.module.resize_token_embeddings([self.model_config.ent_token, self.model_config.sep_token], set_class_token_index = False, From 1f00c9f91e1a6855ff1295ead8ecc75a51ee9209 Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Fri, 14 Feb 2025 11:02:17 +0530 Subject: [PATCH 06/10] Update custom_train.py --- custom_train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/custom_train.py b/custom_train.py index 2a8efe98..23877508 100644 --- a/custom_train.py +++ b/custom_train.py @@ -282,8 +282,8 @@ def train(self, model, optimizer, train_loader, num_steps, device='cuda', rank=N x[k] = v.to(device) try: - with torch.cuda.amp.autocast(dtype=torch.float16): - loss = model(alpha = self.config.loss_alpha, + # with torch.cuda.amp.autocast(dtype=torch.float16): + loss = model(alpha = self.config.loss_alpha, gamma = self.config.loss_gamma, label_smoothing = self.config.label_smoothing, reduction = self.config.loss_reduction, From d5e9c2f9736e369bd471243393f18f7ad5c1eb0b Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Fri, 14 Feb 2025 11:10:40 +0530 Subject: [PATCH 07/10] Update custom_train.py --- custom_train.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/custom_train.py b/custom_train.py index 23877508..a67450a2 100644 --- a/custom_train.py +++ b/custom_train.py @@ -178,12 +178,12 @@ def setup_model_and_optimizer(self, rank=None, device=None): set_class_token_index = False, add_tokens_to_tokenizer=False) if rank is not None: - model = DDP(model, device_ids=[rank], output_device=rank, find_unused_parameters=True) + model = DDP(model, device_ids=[rank], output_device=rank, find_unused_parameters=False) if self.config.labels_encoder is None: model.module.resize_token_embeddings([self.model_config.ent_token, self.model_config.sep_token], set_class_token_index = False, add_tokens_to_tokenizer=False) - optimizer = self.create_optimizer(model.module if isinstance(model, DDP) else model) + optimizer = self.create_optimizer(model.model) if self.compile_model: model.compile_for_training() @@ -194,7 +194,7 @@ def create_dataloader(self, dataset, data_processor, sampler=None, shuffle=True) # dataset = GLiNERDataset(dataset, config = self.config, data_processor=self.data_processor) # collator = DataCollatorWithPadding(self.config) collator = DataCollator(self.config, data_processor=data_processor, prepare_labels=True) - data_loader = DataLoader(dataset, batch_size=self.config.train_batch_size, num_workers=2, + data_loader = DataLoader(dataset, batch_size=self.config.train_batch_size, num_workers=12, shuffle=shuffle, collate_fn=collator, sampler=sampler) return data_loader @@ -208,7 +208,7 @@ def train_dist(self, rank, world_size, dataset): sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True, drop_last=False) - train_loader = train_loader = self.create_dataloader(dataset, model.module.data_processor if isinstance(model, DDP) else model.data_processor, sampler=sampler, shuffle=False) + train_loader = self.create_dataloader(dataset, model.data_processor, sampler=sampler, shuffle=False) num_steps = self.config.num_steps // world_size @@ -282,8 +282,8 @@ def train(self, model, optimizer, train_loader, num_steps, device='cuda', rank=N x[k] = v.to(device) try: - # with torch.cuda.amp.autocast(dtype=torch.float16): - loss = model(alpha = self.config.loss_alpha, + with torch.cuda.amp.autocast(dtype=torch.float16): + loss = model(alpha = self.config.loss_alpha, gamma = self.config.loss_gamma, label_smoothing = self.config.label_smoothing, reduction = self.config.loss_reduction, From 928787b842268ea87bd02abf324bf9f2cc769b9a Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Mon, 24 Feb 2025 22:24:55 +0530 Subject: [PATCH 08/10] Update base.py --- gliner/modeling/base.py | 76 +++++++++++++++++++++++++++++++++-------- 1 file changed, 62 insertions(+), 14 deletions(-) diff --git a/gliner/modeling/base.py b/gliner/modeling/base.py index bcd04cdc..9e4a4b3d 100644 --- a/gliner/modeling/base.py +++ b/gliner/modeling/base.py @@ -15,6 +15,43 @@ from .loss_functions import focal_loss_with_logits from .span_rep import SpanRepLayer +class InterHeadText(nn.Module): + def __init__(self, hidden_dim): + super().__init__() + # Use a single projection for both inputs + self.proj = nn.Linear(hidden_dim, hidden_dim, bias=False) + self.softmax = nn.Softmax(dim=0) + + def forward(self, text1, text2): + # Project both text inputs using the same layer + proj1 = self.proj(text1) + proj2 = self.proj(text2) + + # Compute four interaction scores: + a_11 = (text1 * proj1).sum(-1) # self-interaction of text1 + a_12 = (text1 * proj2).sum(-1) # cross-interaction: text1 with text2's projection + a_21 = (text2 * proj1).sum(-1) # cross-interaction: text2 with text1's projection + a_22 = (text2 * proj2).sum(-1) # self-interaction of text2 + + # Stack and apply softmax to obtain weights + weights1 = self.softmax(torch.stack([a_11, a_12], dim=0)) + weights2 = self.softmax(torch.stack([a_21, a_22], dim=0)) + w11, w12 = weights1.split([1, 1], dim=0) + w21, w22 = weights2.split([1, 1], dim=0) + + # Reshape to allow broadcasting + w11 = w11.squeeze(0).unsqueeze(-1) + w12 = w12.squeeze(0).unsqueeze(-1) + w21 = w21.squeeze(0).unsqueeze(-1) + w22 = w22.squeeze(0).unsqueeze(-1) + + # Create new representations as weighted combinations: + new_text1 = w11 * text1 + w12 * text2 + new_text2 = w21 * text1 + w22 * text2 + + return new_text1, new_text2 + + @dataclass class GLiNERModelOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None @@ -204,13 +241,19 @@ def loss(self, x): class SpanModel(BaseModel): def __init__(self, config, encoder_from_pretrained): super(SpanModel, self).__init__(config, encoder_from_pretrained) - self.span_rep_layer = SpanRepLayer(span_mode = config.span_mode, - hidden_size = config.hidden_size, - max_width = config.max_width, - dropout = config.dropout) - + self.span_rep_layer = SpanRepLayer( + span_mode=config.span_mode, + hidden_size=config.hidden_size, + max_width=config.max_width, + dropout=config.dropout + ) self.prompt_rep_layer = create_projection_layer(config.hidden_size, config.dropout) - + + # Create a stack (n layers) of interaction heads. + self.num_interaction_layers = 8 + self.inter_heads = nn.ModuleList([ + InterHeadText(config.hidden_size) for _ in range(self.num_interaction_layers) + ]) def forward(self, input_ids: Optional[torch.FloatTensor] = None, @@ -229,16 +272,21 @@ def forward(self, labels: Optional[torch.FloatTensor] = None, **kwargs ): - - prompts_embedding, prompts_embedding_mask, words_embedding, mask = self.get_representations(input_ids, attention_mask, - labels_embeddings, labels_input_ids, labels_attention_mask, - text_lengths, words_mask) - span_idx = span_idx*span_mask.unsqueeze(-1) - + # Get initial representations. + prompts_embedding, prompts_embedding_mask, words_embedding, mask = self.get_representations( + input_ids, attention_mask, + labels_embeddings, labels_input_ids, labels_attention_mask, + text_lengths, words_mask + ) + span_idx = span_idx * span_mask.unsqueeze(-1) span_rep = self.span_rep_layer(words_embedding, span_idx) - prompts_embedding = self.prompt_rep_layer(prompts_embedding) - + + # Iteratively apply n interaction layers to enable text-text interaction. + for inter_head in self.inter_heads: + span_rep, prompts_embedding = inter_head(span_rep, prompts_embedding) + + # Compute matching scores via einsum. scores = torch.einsum("BLKD,BCD->BLKC", span_rep, prompts_embedding) loss = None From 023bf7542242ddf69e4e82b2b9dce6f529b325f0 Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Mon, 24 Feb 2025 22:50:06 +0530 Subject: [PATCH 09/10] Update base.py --- gliner/modeling/base.py | 53 ++++++++++++++++++++++++++++++++++------- 1 file changed, 44 insertions(+), 9 deletions(-) diff --git a/gliner/modeling/base.py b/gliner/modeling/base.py index 9e4a4b3d..2db81af4 100644 --- a/gliner/modeling/base.py +++ b/gliner/modeling/base.py @@ -52,6 +52,26 @@ def forward(self, text1, text2): return new_text1, new_text2 +class CrossAttentionHead(nn.Module): + """ + A cross-attention layer that lets span representations (queries) + attend to prompt embeddings (keys/values). This supports different + sequence lengths between the two inputs. + """ + def __init__(self, hidden_dim, num_heads=8, dropout=0.1): + super().__init__() + self.attn = nn.MultiheadAttention(embed_dim=hidden_dim, num_heads=num_heads, batch_first=True, dropout=dropout) + self.norm = nn.LayerNorm(hidden_dim) + self.dropout = nn.Dropout(dropout) + + def forward(self, span_rep, prompt_emb): + # span_rep: (B, N, hidden_dim) where N is total number of spans (flattened) + # prompt_emb: (B, C, hidden_dim) where C is number of prompt tokens + attn_output, _ = self.attn(query=span_rep, key=prompt_emb, value=prompt_emb) + # Apply residual connection and normalization + span_rep = self.norm(span_rep + self.dropout(attn_output)) + return span_rep + @dataclass class GLiNERModelOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None @@ -249,10 +269,11 @@ def __init__(self, config, encoder_from_pretrained): ) self.prompt_rep_layer = create_projection_layer(config.hidden_size, config.dropout) - # Create a stack (n layers) of interaction heads. - self.num_interaction_layers = 8 - self.inter_heads = nn.ModuleList([ - InterHeadText(config.hidden_size) for _ in range(self.num_interaction_layers) + # Create n layers of cross-attention interaction heads. + self.num_interaction_layers = 5 + self.cross_attn_layers = nn.ModuleList([ + CrossAttentionHead(hidden_dim=config.hidden_size, num_heads=8, dropout=config.dropout) + for _ in range(self.num_interaction_layers) ]) def forward(self, @@ -272,21 +293,35 @@ def forward(self, labels: Optional[torch.FloatTensor] = None, **kwargs ): - # Get initial representations. + # Get initial representations (implementation provided elsewhere) prompts_embedding, prompts_embedding_mask, words_embedding, mask = self.get_representations( input_ids, attention_mask, labels_embeddings, labels_input_ids, labels_attention_mask, text_lengths, words_mask ) span_idx = span_idx * span_mask.unsqueeze(-1) + # Obtain span representations: shape (B, L, K, hidden_dim) span_rep = self.span_rep_layer(words_embedding, span_idx) + # Get prompt embeddings: shape (B, C, hidden_dim) prompts_embedding = self.prompt_rep_layer(prompts_embedding) - # Iteratively apply n interaction layers to enable text-text interaction. - for inter_head in self.inter_heads: - span_rep, prompts_embedding = inter_head(span_rep, prompts_embedding) + # The two inputs have different sequence lengths: + # span_rep: (B, L, K, hidden_dim) and prompts_embedding: (B, C, hidden_dim). + # Flatten span_rep over the span dimensions (L and K) to get shape (B, N, hidden_dim), where N = L * K. + B, L, K, D = span_rep.shape + span_rep_flat = span_rep.view(B, L * K, D) + + # Apply each cross-attention layer to update span representations. + for cross_attn in self.cross_attn_layers: + span_rep_flat = cross_attn(span_rep_flat, prompts_embedding) + + # Reshape the updated span representations back to (B, L, K, hidden_dim) + span_rep = span_rep_flat.view(B, L, K, D) - # Compute matching scores via einsum. + # Compute matching scores via einsum: + # span_rep: (B, L, K, hidden_dim) + # prompts_embedding: (B, C, hidden_dim) + # Resulting scores: (B, L, K, C) scores = torch.einsum("BLKD,BCD->BLKC", span_rep, prompts_embedding) loss = None From e017e04117be0c9fc3f128790d544e4b9675e9ff Mon Sep 17 00:00:00 2001 From: Mayank Rakesh <163415854+mayank-rakesh-mck@users.noreply.github.com> Date: Wed, 26 Feb 2025 18:16:09 +0530 Subject: [PATCH 10/10] Update base.py --- gliner/modeling/base.py | 111 +++++----------------------------------- 1 file changed, 14 insertions(+), 97 deletions(-) diff --git a/gliner/modeling/base.py b/gliner/modeling/base.py index 2db81af4..bcd04cdc 100644 --- a/gliner/modeling/base.py +++ b/gliner/modeling/base.py @@ -15,63 +15,6 @@ from .loss_functions import focal_loss_with_logits from .span_rep import SpanRepLayer -class InterHeadText(nn.Module): - def __init__(self, hidden_dim): - super().__init__() - # Use a single projection for both inputs - self.proj = nn.Linear(hidden_dim, hidden_dim, bias=False) - self.softmax = nn.Softmax(dim=0) - - def forward(self, text1, text2): - # Project both text inputs using the same layer - proj1 = self.proj(text1) - proj2 = self.proj(text2) - - # Compute four interaction scores: - a_11 = (text1 * proj1).sum(-1) # self-interaction of text1 - a_12 = (text1 * proj2).sum(-1) # cross-interaction: text1 with text2's projection - a_21 = (text2 * proj1).sum(-1) # cross-interaction: text2 with text1's projection - a_22 = (text2 * proj2).sum(-1) # self-interaction of text2 - - # Stack and apply softmax to obtain weights - weights1 = self.softmax(torch.stack([a_11, a_12], dim=0)) - weights2 = self.softmax(torch.stack([a_21, a_22], dim=0)) - w11, w12 = weights1.split([1, 1], dim=0) - w21, w22 = weights2.split([1, 1], dim=0) - - # Reshape to allow broadcasting - w11 = w11.squeeze(0).unsqueeze(-1) - w12 = w12.squeeze(0).unsqueeze(-1) - w21 = w21.squeeze(0).unsqueeze(-1) - w22 = w22.squeeze(0).unsqueeze(-1) - - # Create new representations as weighted combinations: - new_text1 = w11 * text1 + w12 * text2 - new_text2 = w21 * text1 + w22 * text2 - - return new_text1, new_text2 - - -class CrossAttentionHead(nn.Module): - """ - A cross-attention layer that lets span representations (queries) - attend to prompt embeddings (keys/values). This supports different - sequence lengths between the two inputs. - """ - def __init__(self, hidden_dim, num_heads=8, dropout=0.1): - super().__init__() - self.attn = nn.MultiheadAttention(embed_dim=hidden_dim, num_heads=num_heads, batch_first=True, dropout=dropout) - self.norm = nn.LayerNorm(hidden_dim) - self.dropout = nn.Dropout(dropout) - - def forward(self, span_rep, prompt_emb): - # span_rep: (B, N, hidden_dim) where N is total number of spans (flattened) - # prompt_emb: (B, C, hidden_dim) where C is number of prompt tokens - attn_output, _ = self.attn(query=span_rep, key=prompt_emb, value=prompt_emb) - # Apply residual connection and normalization - span_rep = self.norm(span_rep + self.dropout(attn_output)) - return span_rep - @dataclass class GLiNERModelOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None @@ -261,20 +204,13 @@ def loss(self, x): class SpanModel(BaseModel): def __init__(self, config, encoder_from_pretrained): super(SpanModel, self).__init__(config, encoder_from_pretrained) - self.span_rep_layer = SpanRepLayer( - span_mode=config.span_mode, - hidden_size=config.hidden_size, - max_width=config.max_width, - dropout=config.dropout - ) + self.span_rep_layer = SpanRepLayer(span_mode = config.span_mode, + hidden_size = config.hidden_size, + max_width = config.max_width, + dropout = config.dropout) + self.prompt_rep_layer = create_projection_layer(config.hidden_size, config.dropout) - - # Create n layers of cross-attention interaction heads. - self.num_interaction_layers = 5 - self.cross_attn_layers = nn.ModuleList([ - CrossAttentionHead(hidden_dim=config.hidden_size, num_heads=8, dropout=config.dropout) - for _ in range(self.num_interaction_layers) - ]) + def forward(self, input_ids: Optional[torch.FloatTensor] = None, @@ -293,35 +229,16 @@ def forward(self, labels: Optional[torch.FloatTensor] = None, **kwargs ): - # Get initial representations (implementation provided elsewhere) - prompts_embedding, prompts_embedding_mask, words_embedding, mask = self.get_representations( - input_ids, attention_mask, - labels_embeddings, labels_input_ids, labels_attention_mask, - text_lengths, words_mask - ) - span_idx = span_idx * span_mask.unsqueeze(-1) - # Obtain span representations: shape (B, L, K, hidden_dim) + + prompts_embedding, prompts_embedding_mask, words_embedding, mask = self.get_representations(input_ids, attention_mask, + labels_embeddings, labels_input_ids, labels_attention_mask, + text_lengths, words_mask) + span_idx = span_idx*span_mask.unsqueeze(-1) + span_rep = self.span_rep_layer(words_embedding, span_idx) - # Get prompt embeddings: shape (B, C, hidden_dim) - prompts_embedding = self.prompt_rep_layer(prompts_embedding) - - # The two inputs have different sequence lengths: - # span_rep: (B, L, K, hidden_dim) and prompts_embedding: (B, C, hidden_dim). - # Flatten span_rep over the span dimensions (L and K) to get shape (B, N, hidden_dim), where N = L * K. - B, L, K, D = span_rep.shape - span_rep_flat = span_rep.view(B, L * K, D) - - # Apply each cross-attention layer to update span representations. - for cross_attn in self.cross_attn_layers: - span_rep_flat = cross_attn(span_rep_flat, prompts_embedding) - # Reshape the updated span representations back to (B, L, K, hidden_dim) - span_rep = span_rep_flat.view(B, L, K, D) - - # Compute matching scores via einsum: - # span_rep: (B, L, K, hidden_dim) - # prompts_embedding: (B, C, hidden_dim) - # Resulting scores: (B, L, K, C) + prompts_embedding = self.prompt_rep_layer(prompts_embedding) + scores = torch.einsum("BLKD,BCD->BLKC", span_rep, prompts_embedding) loss = None