This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,3 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from .base import BaseLoss
|
||||
from .mapping import loss_map
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import torch
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from swift.trainers import Trainer, TrainingArguments
|
||||
|
||||
|
||||
class BaseLoss(ABC):
|
||||
"""Abstract base class for custom loss functions.
|
||||
|
||||
This class provides a common interface for implementing custom loss functions
|
||||
that can be integrated with the ms-swift training framework. All custom loss
|
||||
implementations should inherit from this class and implement the __call__ method.
|
||||
|
||||
Attributes:
|
||||
args (TrainingArguments): Training configuration and hyperparameters.
|
||||
trainer (Trainer): Reference to the trainer instance for accessing model
|
||||
and training state.
|
||||
"""
|
||||
|
||||
def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'):
|
||||
"""Initialize the loss function with training arguments and trainer.
|
||||
|
||||
Args:
|
||||
args (TrainingArguments): Training configuration and hyperparameters.
|
||||
trainer (Trainer): Reference to the trainer instance.
|
||||
"""
|
||||
self.args = args
|
||||
self.trainer = trainer
|
||||
mro_class_names = [cls.__name__ for cls in trainer.__class__.__mro__]
|
||||
self.is_megatron = 'BaseMegatronTrainer' in mro_class_names
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, outputs, labels, *, num_items_in_batch=None, loss_scale=None, **kwargs) -> torch.Tensor:
|
||||
"""Calculate the loss value.
|
||||
|
||||
This method must be implemented by all subclasses to define the specific
|
||||
loss calculation logic.
|
||||
|
||||
Args:
|
||||
outputs: Model outputs.
|
||||
labels: Ground truth labels or targets.
|
||||
num_items_in_batch (int, optional): Number of items (tokens) in the current batch,
|
||||
Defaults to None.
|
||||
loss_scale (float, optional): Scaling factor to apply to the loss value.
|
||||
Defaults to None.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A scalar tensor representing the computed loss value.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,12 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from .base import BaseLoss
|
||||
|
||||
|
||||
class CustomCrossEntropyLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, *, num_items_in_batch=None, loss_scale=None, **kwargs):
|
||||
from swift.trainers import per_token_loss_func
|
||||
token_loss = per_token_loss_func(outputs, labels)
|
||||
if num_items_in_batch is None:
|
||||
num_items_in_batch = (labels[:, 1:] != -100).sum()
|
||||
return token_loss.sum() / num_items_in_batch
|
||||
@@ -0,0 +1,321 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from accelerate.utils import gather_object
|
||||
from enum import Enum
|
||||
from torch import nn
|
||||
from torch.nn import MSELoss
|
||||
from transformers.utils import strtobool
|
||||
|
||||
from swift.sequence_parallel import sequence_parallel
|
||||
from swift.utils import get_dist_setting
|
||||
from .base import BaseLoss
|
||||
|
||||
|
||||
# Code borrowed from sentence_transformers
|
||||
class SiameseDistanceMetric(Enum):
|
||||
"""The metric for the contrastive loss"""
|
||||
|
||||
EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2) # noqa
|
||||
MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1) # noqa
|
||||
COSINE_DISTANCE = lambda x, y: 1 - F.cosine_similarity(x, y) # noqa
|
||||
|
||||
|
||||
def _parse_pair_sentence(outputs):
|
||||
if isinstance(outputs, dict):
|
||||
last_hidden_state = outputs['last_hidden_state']
|
||||
else:
|
||||
last_hidden_state = outputs
|
||||
batch_size = last_hidden_state.shape[0]
|
||||
shape_len = len(last_hidden_state.shape)
|
||||
first_sentence = list(range(0, batch_size, 2))
|
||||
second_sentence = list(range(1, batch_size, 2))
|
||||
if shape_len == 3:
|
||||
sentence1 = last_hidden_state[first_sentence][:, 0].squeeze(dim=1)
|
||||
sentence2 = last_hidden_state[second_sentence][:, 0].squeeze(dim=1)
|
||||
else:
|
||||
sentence1 = last_hidden_state[first_sentence]
|
||||
sentence2 = last_hidden_state[second_sentence]
|
||||
return sentence1, sentence2
|
||||
|
||||
|
||||
class CosineSimilarityLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
|
||||
# You need to return a scalar representing the loss.
|
||||
cos_score_transformation = nn.Identity()
|
||||
loss_fct = MSELoss()
|
||||
sentence1, sentence2 = _parse_pair_sentence(outputs)
|
||||
output = cos_score_transformation(torch.cosine_similarity(sentence1, sentence2))
|
||||
return loss_fct(output, labels.to(output.dtype).view(-1))
|
||||
|
||||
|
||||
class ContrastiveLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
|
||||
sentence1, sentence2 = _parse_pair_sentence(outputs)
|
||||
distance_metric = SiameseDistanceMetric.COSINE_DISTANCE
|
||||
distances = distance_metric(sentence1, sentence2)
|
||||
margin = 0.5
|
||||
labels = labels.to(sentence1.dtype)
|
||||
losses = 0.5 * (labels * distances.pow(2) + (1 - labels) * F.relu(margin - distances).pow(2))
|
||||
return losses.mean()
|
||||
|
||||
|
||||
class OnlineContrastiveLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
|
||||
sentence1, sentence2 = _parse_pair_sentence(outputs)
|
||||
distance_metric = SiameseDistanceMetric.COSINE_DISTANCE
|
||||
distance_matrix = distance_metric(sentence1, sentence2)
|
||||
negs = distance_matrix[labels == 0]
|
||||
poss = distance_matrix[labels == 1]
|
||||
|
||||
# select hard positive and hard negative pairs
|
||||
negative_pairs = negs[negs < (poss.max() if len(poss) > 1 else negs.mean())]
|
||||
positive_pairs = poss[poss > (negs.min() if len(negs) > 1 else poss.mean())]
|
||||
|
||||
positive_loss = positive_pairs.pow(2).sum()
|
||||
margin = 0.5
|
||||
negative_loss = F.relu(margin - negative_pairs).pow(2).sum()
|
||||
loss = positive_loss + negative_loss
|
||||
return loss
|
||||
|
||||
|
||||
def _parse_multi_negative_sentences(sentences, labels, hard_negatives=None):
|
||||
split_indices = torch.nonzero(labels, as_tuple=False).squeeze().tolist()
|
||||
if isinstance(split_indices, int):
|
||||
split_indices = [split_indices]
|
||||
split_indices.append(len(labels))
|
||||
split_indices = np.array(split_indices) + np.array(list(range(len(split_indices))))
|
||||
split_tensors = []
|
||||
|
||||
for i in range(len(split_indices) - 1):
|
||||
start = split_indices[i]
|
||||
end = split_indices[i + 1]
|
||||
split_part = sentences[start:end]
|
||||
if hard_negatives is not None:
|
||||
negatives = len(split_part) - 2
|
||||
assert negatives > 0
|
||||
if negatives > hard_negatives:
|
||||
split_part = split_part[:hard_negatives + 2]
|
||||
elif negatives < hard_negatives:
|
||||
selected = np.random.choice(list(range(negatives)), size=hard_negatives - negatives, replace=True)
|
||||
selected += 1 # skip positive
|
||||
split_part = torch.cat((split_part, split_part[selected]), dim=0)
|
||||
split_tensors.append(split_part)
|
||||
return split_tensors
|
||||
|
||||
|
||||
class InfonceLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
|
||||
temperature = float(os.environ.get('INFONCE_TEMPERATURE', '0.1')) # temperature
|
||||
# calculate CE across the batch, meaning all samples will be negative except the matching positive
|
||||
use_batch = strtobool(os.environ.get('INFONCE_USE_BATCH', 'True'))
|
||||
hard_negatives = os.environ.get('INFONCE_HARD_NEGATIVES', None) # how many negative prompts kept in one sample
|
||||
# mask out fake negatives
|
||||
infonce_mask_fake_negative = strtobool(os.environ.get('INFONCE_MASK_FAKE_NEGATIVE', 'False'))
|
||||
fake_neg_margin = float(os.environ.get('INFONCE_FAKE_NEG_MARGIN', '0.1'))
|
||||
# enhanced components to align with Qwen3-Embedding denominator; controlled individually
|
||||
# defaults set to False for backward compatibility
|
||||
infonce_include_qq = strtobool(os.environ.get('INFONCE_INCLUDE_QQ', 'False'))
|
||||
infonce_include_dd = strtobool(os.environ.get('INFONCE_INCLUDE_DD', 'False'))
|
||||
if hard_negatives is not None:
|
||||
hard_negatives = int(hard_negatives)
|
||||
if self.is_megatron:
|
||||
from megatron.core import mpu
|
||||
rank, world_size = mpu.get_data_parallel_rank(), mpu.get_data_parallel_world_size()
|
||||
else:
|
||||
rank, _, world_size, _ = get_dist_setting()
|
||||
# repeat of anchor(1)+positive(1)+negatives(n)
|
||||
sentences = outputs['last_hidden_state']
|
||||
|
||||
if world_size > 1 and use_batch:
|
||||
if getattr(sequence_parallel, 'dp_group', None) is not None:
|
||||
all_sentences = sequence_parallel._gather_object_dp(sentences.unsqueeze(0))
|
||||
labels = sequence_parallel._gather_object_dp(labels)
|
||||
rank = sequence_parallel.dp_rank
|
||||
elif self.is_megatron:
|
||||
from megatron.core import mpu
|
||||
dp_group = mpu.get_data_parallel_group()
|
||||
shapes = [sentences.new_empty((2, ), dtype=torch.long) for _ in range(world_size)]
|
||||
dist.all_gather(
|
||||
shapes,
|
||||
sentences.new_tensor(sentences.shape, dtype=torch.long),
|
||||
group=dp_group,
|
||||
)
|
||||
all_sentences = [sentences.new_empty(shape.tolist()) for shape in shapes]
|
||||
dist.all_gather(
|
||||
all_sentences,
|
||||
sentences,
|
||||
group=dp_group,
|
||||
)
|
||||
else:
|
||||
# gather all the sentences and labels across the gpus when calculate loss across all batches of all gpus
|
||||
all_sentences = gather_object(sentences.unsqueeze(0))
|
||||
labels = gather_object(labels)
|
||||
# override the gathered one
|
||||
all_sentences[rank] = sentences
|
||||
for idx in range(len(all_sentences)):
|
||||
if idx == rank:
|
||||
continue
|
||||
# we don't calculate grad from other gpus
|
||||
all_sentences[idx] = all_sentences[idx].detach().to(sentences.device)
|
||||
sentences = torch.cat(all_sentences, dim=0)
|
||||
labels = [tensor.to(sentences.device) for tensor in labels]
|
||||
labels = torch.stack(labels, dim=0)
|
||||
|
||||
# split tensors into single sample
|
||||
# for example: batch_size=2 with tensor anchor(1)+positive(1)+negatives(3) + anchor(1)+positive(1)+negatives(2)
|
||||
# labels will be [1,0,0,0,1,0,0], meaning 1 positive, 3 negatives, 1 positive, 2 negatives
|
||||
split_tensors = _parse_multi_negative_sentences(sentences, labels, hard_negatives)
|
||||
loss = 0
|
||||
can_batched = hard_negatives is not None
|
||||
if hard_negatives is None and len(set([s.shape[0] for s in split_tensors])) == 1:
|
||||
# all tensors have the same batch size
|
||||
can_batched = True
|
||||
if not use_batch:
|
||||
# only calculate loss inside one sample
|
||||
if can_batched:
|
||||
# negative numbers are equal
|
||||
# [B, neg+2, D]
|
||||
sentences = torch.stack(split_tensors, dim=0)
|
||||
# [B, 1, D] * [B, neg+1, D]
|
||||
similarity_matrix = torch.matmul(sentences[:, 0:1], sentences[:, 1:].transpose(1, 2)) / temperature
|
||||
# The positive one is the first element
|
||||
labels = torch.zeros(len(split_tensors), dtype=torch.int64).to(sentences.device)
|
||||
loss = nn.CrossEntropyLoss()(similarity_matrix.squeeze(1), labels)
|
||||
else:
|
||||
# the negative numbers may be different, use for loop
|
||||
for tensor in split_tensors:
|
||||
# [D] * [neg+1, D]
|
||||
similarity_matrix = torch.matmul(tensor[0], tensor[1:].T) / temperature
|
||||
# The positive one is the first element
|
||||
labels = torch.tensor(0).to(tensor.device)
|
||||
loss += nn.CrossEntropyLoss()(similarity_matrix, labels)
|
||||
# avg between all batches in one gpu
|
||||
loss /= len(split_tensors)
|
||||
else:
|
||||
if can_batched:
|
||||
# [B, neg+2, D]
|
||||
sentences = torch.stack(split_tensors, dim=0)
|
||||
# base q->d similarities (includes own positive and all in-batch documents)
|
||||
queries = sentences[:, 0].squeeze(1) # [B, D]
|
||||
docs_all = sentences[:, 1:].reshape(-1, sentences.size(2)) # [B*(neg+1), D]
|
||||
qd_matrix = torch.matmul(queries, docs_all.T) # [B, B*(neg+1)]
|
||||
# target indices: start of each group's document block (its positive)
|
||||
labels = torch.tensor(range(0,
|
||||
sentences.size(0) * (sentences.size(1) - 1),
|
||||
sentences.size(1) - 1)).view(-1).to(sentences.device)
|
||||
|
||||
logits_list = [qd_matrix]
|
||||
|
||||
if infonce_include_qq:
|
||||
# q->q similarities; exclude self via -inf on diagonal to avoid accidental positives
|
||||
qq_matrix = torch.matmul(queries, queries.T) # [B, B]
|
||||
qq_matrix = qq_matrix.clone()
|
||||
qq_matrix.fill_diagonal_(float('-inf'))
|
||||
logits_list.append(qq_matrix)
|
||||
|
||||
if infonce_include_dd:
|
||||
# d+ -> d (doc-doc) similarities; exclude self-positive column per row
|
||||
pos_docs = sentences[:, 1].squeeze(1) # [B, D]
|
||||
dd_matrix = torch.matmul(pos_docs, docs_all.T) # [B, B*(neg+1)]
|
||||
# mask self positive per row: column index = row_idx * (neg+1)
|
||||
block = sentences.size(1) - 1 # (neg+1)
|
||||
if block > 0:
|
||||
row_idx = torch.arange(dd_matrix.size(0), device=dd_matrix.device)
|
||||
col_idx = row_idx * block
|
||||
dd_matrix[row_idx, col_idx] = float('-inf')
|
||||
logits_list.append(dd_matrix)
|
||||
|
||||
if infonce_mask_fake_negative:
|
||||
# thresholds derived from positive q->d scores per row
|
||||
row_idx = torch.arange(qd_matrix.size(0), device=qd_matrix.device)
|
||||
pos_scores = qd_matrix[row_idx, labels]
|
||||
thresholds = pos_scores.view(-1, 1).detach() + fake_neg_margin
|
||||
|
||||
# qd block mask
|
||||
qd_block = qd_matrix.clone()
|
||||
qd_mask = qd_block > thresholds
|
||||
qd_block[qd_mask] = float('-inf')
|
||||
|
||||
components = [qd_block]
|
||||
|
||||
# qq block mask (if present)
|
||||
if infonce_include_qq:
|
||||
qq_block = qq_matrix.clone()
|
||||
qq_mask = qq_block > thresholds
|
||||
qq_block[qq_mask] = float('-inf')
|
||||
# diagonal already masked unconditionally at construction time
|
||||
components.append(qq_block)
|
||||
|
||||
# dd block (if present): self-positive column already masked unconditionally
|
||||
if infonce_include_dd:
|
||||
# align with Qwen3-Embedding, no threshold masking for d-d
|
||||
components.append(dd_matrix)
|
||||
|
||||
similarity_matrix = torch.cat(components, dim=1)
|
||||
else:
|
||||
# concatenate all components without masking
|
||||
similarity_matrix = torch.cat(logits_list, dim=1)
|
||||
# temperature scaling and CE
|
||||
similarity_matrix = similarity_matrix / temperature
|
||||
loss = nn.CrossEntropyLoss()(similarity_matrix, labels)
|
||||
else:
|
||||
all_tensors = []
|
||||
for tensor in split_tensors:
|
||||
all_tensors.append(tensor[1:])
|
||||
# cat all neg+1 tensors
|
||||
sentences = torch.cat(all_tensors, dim=0)
|
||||
# prepare query anchors list if q-q is included
|
||||
if infonce_include_qq:
|
||||
queries_all = torch.stack([t[0] for t in split_tensors], dim=0) # [B, D]
|
||||
length = 0
|
||||
for idx, tensor in enumerate(split_tensors):
|
||||
# [D] * [B*(neg+1), D], neg numbers are different
|
||||
qd_vec = torch.matmul(tensor[0], sentences.T)
|
||||
target = torch.tensor(length).to(tensor.device)
|
||||
logits_parts = []
|
||||
|
||||
# compute threshold from positive q->d score
|
||||
threshold = (qd_vec[target].detach() + fake_neg_margin)
|
||||
|
||||
# qd part with masking
|
||||
if infonce_mask_fake_negative:
|
||||
qd_masked = torch.where(qd_vec > threshold, torch.tensor(float('-inf'), device=qd_vec.device),
|
||||
qd_vec)
|
||||
else:
|
||||
qd_masked = qd_vec
|
||||
logits_parts.append(qd_masked)
|
||||
|
||||
# qq part
|
||||
if infonce_include_qq:
|
||||
qq_vec = torch.matmul(tensor[0], queries_all.T) # [B]
|
||||
# exclude self
|
||||
qq_vec = qq_vec.clone()
|
||||
qq_vec[idx] = float('-inf')
|
||||
if infonce_mask_fake_negative:
|
||||
qq_vec = torch.where(qq_vec > threshold, torch.tensor(float('-inf'), device=qq_vec.device),
|
||||
qq_vec)
|
||||
logits_parts.append(qq_vec)
|
||||
|
||||
# dd part
|
||||
if infonce_include_dd:
|
||||
dd_vec = torch.matmul(tensor[1], sentences.T) # [B*(neg+1)]
|
||||
# mask self positive column for this row only (no threshold masking for d-d)
|
||||
block = split_tensors[idx].size(0) - 1 # (neg+1) for this group
|
||||
dd_vec[length] = float('-inf')
|
||||
logits_parts.append(dd_vec)
|
||||
|
||||
logits_row = torch.cat(logits_parts, dim=-1)
|
||||
logits_row = logits_row / temperature
|
||||
loss += nn.CrossEntropyLoss()(logits_row.unsqueeze(0), target.unsqueeze(0))
|
||||
# next positive is neg+1
|
||||
length += tensor.size(0) - 1
|
||||
loss /= len(split_tensors)
|
||||
return loss
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from .causal_lm import CustomCrossEntropyLoss
|
||||
from .embedding import ContrastiveLoss, CosineSimilarityLoss, InfonceLoss, OnlineContrastiveLoss
|
||||
from .reranker import ListwiseRerankerLoss, PointwiseRerankerLoss
|
||||
|
||||
loss_map = {
|
||||
'cross_entropy': CustomCrossEntropyLoss, # examples
|
||||
# embedding
|
||||
'cosine_similarity': CosineSimilarityLoss,
|
||||
'contrastive': ContrastiveLoss,
|
||||
'online_contrastive': OnlineContrastiveLoss,
|
||||
'infonce': InfonceLoss,
|
||||
# # reranker
|
||||
'pointwise_reranker': PointwiseRerankerLoss,
|
||||
'listwise_reranker': ListwiseRerankerLoss,
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from torch.nn import CrossEntropyLoss
|
||||
|
||||
from swift.utils import get_last_valid_indices
|
||||
from .base import BaseLoss
|
||||
|
||||
|
||||
class PointwiseRerankerLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, **kwargs) -> torch.Tensor:
|
||||
logits = outputs.logits
|
||||
logits = logits.squeeze(1)
|
||||
labels = labels.to(logits.dtype)
|
||||
loss_fct = nn.BCEWithLogitsLoss()
|
||||
loss = loss_fct(logits, labels)
|
||||
return loss
|
||||
|
||||
|
||||
class ListwiseRerankerLoss(BaseLoss):
|
||||
|
||||
def __call__(self, outputs, labels, **kwargs):
|
||||
"""
|
||||
List-wise reranker loss function.
|
||||
|
||||
This loss function groups samples by query based on the pattern where each group
|
||||
consists of 1 positive document followed by n negative documents. It treats the
|
||||
ranking task as a classification problem within each group, using cross-entropy
|
||||
loss to identify the positive document among all candidates.
|
||||
|
||||
Data format expected:
|
||||
- labels: [1, 0, 0, 0, 1, 0, 0, ...] where 1 indicates positive, 0 indicates negative
|
||||
- Each 1 is followed by its corresponding negative documents until the next 1
|
||||
|
||||
Environment variables for configuration:
|
||||
- LISTWISE_RERANKER_TEMPERATURE: Temperature for softmax (default: 1.0)
|
||||
- LISTWISE_RERANKER_MIN_GROUP_SIZE: Minimum group size to include (default: 2)
|
||||
|
||||
Args:
|
||||
outputs: Model outputs containing logits [batch_size, 1]
|
||||
labels: Binary labels (1 for positive, 0 for negative) [batch_size]
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Cross entropy loss for ranking classification
|
||||
"""
|
||||
logits = outputs.logits.squeeze(-1) # [batch_size]
|
||||
labels = labels.float()
|
||||
|
||||
# Configuration from environment variables
|
||||
temperature = float(os.environ.get('LISTWISE_RERANKER_TEMPERATURE', '1.0'))
|
||||
min_group_size = int(os.environ.get('LISTWISE_RERANKER_MIN_GROUP_SIZE', '2'))
|
||||
|
||||
# Find positive sample indices to determine group boundaries
|
||||
positive_indices = torch.nonzero(labels == 1, as_tuple=False).squeeze(-1)
|
||||
|
||||
if len(positive_indices) == 0:
|
||||
# No positive samples in this batch, return zero loss
|
||||
return torch.tensor(0.0, device=logits.device, requires_grad=True)
|
||||
|
||||
# Ensure positive_indices is 1D
|
||||
if positive_indices.dim() == 0:
|
||||
positive_indices = positive_indices.unsqueeze(0)
|
||||
|
||||
total_loss = 0.0
|
||||
num_groups = 0
|
||||
|
||||
for i, pos_idx in enumerate(positive_indices):
|
||||
# Determine group boundaries
|
||||
group_start = pos_idx.item()
|
||||
|
||||
# Find the end of current group (start of next group or end of batch)
|
||||
if i + 1 < len(positive_indices):
|
||||
group_end = positive_indices[i + 1].item()
|
||||
else:
|
||||
group_end = len(labels)
|
||||
|
||||
# Extract group logits and labels
|
||||
group_logits = logits[group_start:group_end] # [group_size]
|
||||
group_labels = labels[group_start:group_end] # [group_size]
|
||||
|
||||
# Skip groups that are too small
|
||||
if len(group_logits) < min_group_size:
|
||||
continue
|
||||
|
||||
# Verify that the first sample in the group is positive
|
||||
if group_labels[0] != 1:
|
||||
continue # Skip malformed groups
|
||||
|
||||
# Apply temperature scaling for better training dynamics
|
||||
scaled_logits = group_logits / temperature
|
||||
|
||||
# The positive document is always at index 0 within the group
|
||||
target = torch.tensor(0, dtype=torch.long, device=logits.device)
|
||||
|
||||
# Apply cross-entropy loss: positive document should have highest score
|
||||
loss_fct = CrossEntropyLoss()
|
||||
group_loss = loss_fct(scaled_logits.unsqueeze(0), target.unsqueeze(0))
|
||||
|
||||
total_loss += group_loss
|
||||
num_groups += 1
|
||||
|
||||
if num_groups == 0:
|
||||
return torch.tensor(0.0, device=logits.device, requires_grad=True)
|
||||
|
||||
# Return average loss across all groups
|
||||
return total_loss / num_groups
|
||||
Reference in New Issue
Block a user