# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import math
from dataclasses import dataclass, field
import torch
from fairseq import metrics, utils
from fairseq.criterions import FairseqCriterion, register_criterion
from fairseq.criterions.label_smoothed_cross_entropy import label_smoothed_nll_loss
from fairseq.dataclass import FairseqDataclass
from omegaconf import II
# 这段代码是用 Python 的 dataclass 装饰器定义了一个名为 LabelSmoothedCrossEntropyCriterionWithSimCutConfig 的数据类，用于存储模型训练过程中的配置参数。
@dataclass
class LabelSmoothedCrossEntropyCriterionWithSimCutConfig(FairseqDataclass):
    label_smoothing: float = field(
      # 表示平滑系数，为0则表示没有平滑。
        default=0.0,
        metadata={"help": "epsilon for label smoothing, 0 means no label smoothing"},
    )
    alpha: float = field(
      #  表示α参数，为0则表示不进行simcut  (正则化系数)
        default=0.0,
        metadata={"help": "alpha for simcut, 0 means no simcut"},
    )
    p: float = field(
      # cutoff的概率参数，为0则表示不进行simcut中的cutoff操作
        default=0.0,
        metadata={"help": "probability for cutoff in simcut, 0 means no cutoff in simcut"},
    )
    report_accuracy: bool = field(
      # 表示是否需要报告准确性指标。
        default=False,
        metadata={"help": "report accuracy metric"},
    )
    ignore_prefix_size: int = field(
      # 忽略token的数量。
        default=0,
        metadata={"help": "Ignore first N tokens"},
    )
    # 布尔值，即是否进行句子级别的平均，具体作用是用于计算loss值时，是否对每个句子的loss求平均，如果为True，则对整个batch的所有句子的loss值求平均。
    sentence_avg: bool = II("optimization.sentence_avg")
# 这段代码是使用fairseq库中的@register_criterion装饰器将LabelSmoothedCrossEntropyCriterionWithSimCutConfig类注册为label_smoothed_cross_entropy_with_simcut
# 这意味着在训练模型时可以使用label_smoothed_cross_entropy_with_simcut作为损失函数。
# @register_criterion装饰器提供了一种方便的方式，使用户可以自定义和注册新的损失函数。
@register_criterion(
    "label_smoothed_cross_entropy_with_simcut", dataclass=LabelSmoothedCrossEntropyCriterionWithSimCutConfig
)
# 该类实现了'simcut'方法，用于计算相似度裁剪损失。
class LabelSmoothedCrossEntropyWithSimCutCriterion(FairseqCriterion):
    def __init__(
      # 该类接受一下参数
        self,
        # 任务对象
        task,
        # 一个布尔值，表示损失是否以句子为单位进行平均。如果为True，则每个句子的损失将被平均，否则将对所有标记的损失进行平均。
        sentence_avg,
        # 平滑
        label_smoothing,
        # 正则化项系数
        alpha=0.0,
        # 表示simcut方法中的p值
        p=0.0,
        # 一个整数，表示需要忽略的前缀的长度。
        ignore_prefix_size=0,
        # 是否需要报告准确性指标
        report_accuracy=False,
    ):
        super().__init__(task)
        self.sentence_avg = sentence_avg
        self.eps = label_smoothing
        self.ignore_prefix_size = ignore_prefix_size
        self.report_accuracy = report_accuracy
        self.alpha = alpha
        self.p = 1 - p
    def simcut(self, model, omega, sample, reduce):
      """
    这段代码定义了一个名为simcut的方法，接受四个参数：self，model，omega，sample和reduce。
    """
    # 计算omega矩阵在最后一维的softmax，得到一个概率分布，保存在变量prob中。
        prob = torch.nn.functional.softmax(omega, dim=-1)
    # 判断target序列中是否有padding字符，并生成一个与target同样形状的张量，其中为True的位置代表
    # 原来target中非padding字符的位置。
        valid_indices = (sample["target"] != self.padding_idx)
    # 调用模型中的encoder模块对输入进行编码，生成编码器输出。
        encoder_out = model.encoder(
            src_tokens=sample["net_input"]["src_tokens"],
            src_lengths=sample["net_input"]["src_lengths"],
            simcut_p=self.p)
        # 调用模型中的decoder模块对输入进行解码，生成解码器输出。 
        decoder_out = model.decoder(
            prev_output_tokens=sample["net_input"]["prev_output_tokens"],
            encoder_out=encoder_out,
            simcut_p=self.p)
    # 使用KL散度来计算解码器输出分布与目标分布之间的差异，并将结果保存在loss变量之中。
        loss = torch.nn.functional.kl_div(
            input=torch.nn.functional.log_softmax(decoder_out[0], dim=-1),
            target=prob, reduction='none')
        # 将每个位置的KL散度相加，得到总的散度。
        loss = loss.sum(dim=-1)
        # 将padding位置的kl散度值置为0。
        loss = loss * valid_indices.float()
        # 
        if reduce:
            loss = loss.sum()
        return loss
    def forward(self, model, sample, reduce=True):
        """Compute the loss for the given sample.
        用于计算给定样本的损失。
        Returns a tuple with three elements:
        1) the loss
        2) the sample size, which is used as the denominator for the gradient
        3) logging outputs to display while training
        """
        # 调用模型对输入进行前向传播，得到模型的输出
        net_output = model(**sample["net_input"])
        # 计算模型输出的损失，同时也计算了负对数似然损失，并将结果保存在loss和nll_loss中去。
        loss, nll_loss = self.compute_loss(model, net_output, sample, reduce=reduce)
    # 
        if model.training:
          # 调用simcut方法，计算基于KL散度的相似度正则项，并将其加到总的损失函数中去。
            loss += self.alpha * self.simcut(model, net_output[0], sample, reduce)
    # 计算样本的大小，用于作为梯度的分母。
        sample_size = (
            sample["target"].size(0) if self.sentence_avg else sample["ntokens"]
        )
        # 将损失函数的值、负对数似然损失的值、样本中的token数，句子数和样本大小等信息保存在一个字典中。
        logging_output = {
            "loss": loss.data,
            "nll_loss": nll_loss.data,
            "ntokens": sample["ntokens"],
            "nsentences": sample["target"].size(0),
            "sample_size": sample_size,
        }
        # 如果设置了报告准确率的标志，则执行下面的操作。
        if self.report_accuracy:
          # 计算模型的准确率，并将正确预测的数量和总数分别保存在n_correct和total变量中。
            n_correct, total = self.compute_accuracy(model, net_output, sample)
            logging_output["n_correct"] = utils.item(n_correct.data)
            logging_output["total"] = utils.item(total.data)
        return loss, sample_size, logging_output
    def get_lprobs_and_target(self, model, net_output, sample):
      # 获取归一化的对数概率，net_output代表输出，log_probs代表需要返回对数概率而非原始概率。
        lprobs = model.get_normalized_probs(net_output, log_probs=True)
        # 获取目标序列，其中sample是模型的输入，net_output是输出。
        target = model.get_targets(sample, net_output)
        # 如果ignore_prefix_size 大于0，则会忽略输出序列的前缀部分。
        if self.ignore_prefix_size > 0:
            # lprobs: B x T x C
            lprobs = lprobs[:, self.ignore_prefix_size :, :].contiguous()
            target = target[:, self.ignore_prefix_size :].contiguous()
        # 函数返回对数概率和目标序列，
        return lprobs.view(-1, lprobs.size(-1)), target.view(-1)
    def compute_loss(self, model, net_output, sample, reduce=True):
      """
    用于计算给定模型在某个输入样本上的损失值。
    """
    # self 是指代类实例本身，model 是正在评估的模型，net_output 是给定输入样本时模型的输出，sample 是输入样本本身。
    # 用于计算预测对数概率值和目标值
        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
        # 用label_smoothed_nll_loss函数来计算损失值，该函数返回两个值，总损失值loss和负对数似然值。
        loss, nll_loss = label_smoothed_nll_loss(
            lprobs,
            target,
            self.eps,
            ignore_index=self.padding_idx,
            # reduce 是一个布尔型标志，用于指示损失是否应该被缩减为一个标量值（例如，是否应该在批次上进行平均）。
            reduce=reduce,
        )
        return loss, nll_loss
    def compute_accuracy(self, model, net_output, sample):
      """
    计算神经网络模型在给定输入样本上预测准确率的方法。
    """
    # 
        lprobs, target = self.get_lprobs_and_target(model, net_output, sample)
        # 通过target.ne(self.padding_idx)创建一个掩码张量。这会创建一个二进制掩码张量，其中每个元素为1，如果target中对应的元素不等于填充索引。
        mask = target.ne(self.padding_idx)
        n_correct = torch.sum(
            lprobs.argmax(1).masked_select(mask).eq(target.masked_select(mask))
        )
        total = torch.sum(mask)
        return n_correct, total
    @classmethod
    def reduce_metrics(cls, logging_outputs) -> None:
        """Aggregate logging outputs from data parallel training."""
        # 用于在数据并行训练时聚合日志输出。
        # logging_outputs：这个参数包含了多个日志输出的列表
    # 循环遍历logging_outputs列表，对于每个日志输出，统计它的loss，nll_loss、ntokens和sample_size指标的总和。
        loss_sum = sum(log.get("loss", 0) for log in logging_outputs)
        nll_loss_sum = sum(log.get("nll_loss", 0) for log in logging_outputs)
        ntokens = sum(log.get("ntokens", 0) for log in logging_outputs)
        sample_size = sum(log.get("sample_size", 0) for log in logging_outputs)
    # 调用metrics.log_scalar和log_derived方法，将聚合后的指标值记录下来，其中log_scalar方法记录标量指标，如loss和nll_loss, log_derived()方法记录的是可以从已记录的指标派生出来的新指标，如ppl和accuracy。
        metrics.log_scalar(
            "loss", loss_sum / sample_size / math.log(2), sample_size, round=3
        )
        metrics.log_scalar(
            "nll_loss", nll_loss_sum / ntokens / math.log(2), ntokens, round=3
        )
        metrics.log_derived(
            "ppl", lambda meters: utils.get_perplexity(meters["nll_loss"].avg)
        )
        total = utils.item(sum(log.get("total", 0) for log in logging_outputs))
        if total > 0:
            metrics.log_scalar("total", total)
            n_correct = utils.item(
                sum(log.get("n_correct", 0) for log in logging_outputs)
            )
            metrics.log_scalar("n_correct", n_correct)
            metrics.log_derived(
                "accuracy",
                lambda meters: round(
                    meters["n_correct"].sum * 100.0 / meters["total"].sum, 3
                )
                if meters["total"].sum > 0
                else float("nan"),
            )
    @staticmethod
    def logging_outputs_can_be_summed() -> bool:
        """
        Whether the logging outputs returned by `forward` can be summed
        across workers prior to calling `reduce_metrics`. Setting this
        to True will improves distributed training speed.
        """
        return True

absl-py==1.0.0
antlr4-python3-runtime==4.8
anyio==3.4.0
argon2-cffi==21.1.0
attrs==21.2.0
Babel==2.9.1
backcall==0.2.0
bitarray==2.7.3
bleach==4.1.0
brotlipy==0.7.0
cachetools==4.2.4
certifi==2021.5.30
cffi @ file:///tmp/build/80754af9/cffi_1625807838443/work
chardet @ file:///tmp/build/80754af9/chardet_1607706746162/work
cmake==3.26.0
colorama==0.4.6
conda==4.10.3
conda-package-handling @ file:///tmp/build/80754af9/conda-package-handling_1618262148928/work
cryptography @ file:///tmp/build/80754af9/cryptography_1616769286105/work
cycler==0.11.0
Cython==0.29.33
debugpy==1.5.1
decorator==5.1.0
defusedxml==0.7.1
entrypoints==0.3
-e git+https://github.com/gpengzhi/Bi-SimCut@ab38d206413f32da45b6c0e0f9ee7f15d0d48a0f#egg=fairseq&subdirectory=fairseq
filelock==3.10.0
fonttools==4.28.2
google-auth==2.3.3
google-auth-oauthlib==0.4.6
grpcio==1.42.0
hydra-core==1.0.7
idna @ file:///home/linux1/recipes/ci/idna_1610986105248/work
importlib-metadata==4.8.2
importlib-resources==5.4.0
ipykernel==6.5.1
ipython==7.29.0
ipython-genutils==0.2.0
ipywidgets==7.6.5
jedi==0.18.1
Jinja2==3.0.3
json5==0.9.6
jsonschema==4.2.1
jupyter-client==7.1.0
jupyter-core==4.9.1
jupyter-server==1.12.0
jupyterlab==3.2.4
jupyterlab-language-pack-zh-CN==3.2.post2
jupyterlab-pygments==0.1.2
jupyterlab-server==2.8.2
jupyterlab-widgets==1.0.2
kiwisolver==1.3.2
lightseq==3.0.1
lit==15.0.7
lxml==4.9.2
Markdown==3.3.6
MarkupSafe==2.0.1
matplotlib==3.5.0
matplotlib-inline==0.1.3
mistune==0.8.4
mpmath==1.3.0
nbclassic==0.3.4
nbclient==0.5.9
nbconvert==6.3.0
nbformat==5.1.3
nest-asyncio==1.5.1
networkx==3.0
ninja==1.11.1
notebook==6.4.6
numpy==1.21.4
nvidia-cublas-cu11==11.10.3.66
nvidia-cuda-cupti-cu11==11.7.101
nvidia-cuda-nvrtc-cu11==11.7.99
nvidia-cuda-runtime-cu11==11.7.99
nvidia-cudnn-cu11==8.5.0.96
nvidia-cufft-cu11==10.9.0.58
nvidia-curand-cu11==10.2.10.91
nvidia-cusolver-cu11==11.4.0.1
nvidia-cusparse-cu11==11.7.4.91
nvidia-nccl-cu11==2.14.3
nvidia-nvtx-cu11==11.7.91
oauthlib==3.1.1
omegaconf==2.0.6
packaging==21.3
pandocfilters==1.5.0
parso==0.8.2
pexpect==4.8.0
pickleshare==0.7.5
Pillow==8.4.0
portalocker==2.7.0
prometheus-client==0.12.0
prompt-toolkit==3.0.22
protobuf==3.19.1
ptyprocess==0.7.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pycosat==0.6.3
pycparser @ file:///tmp/build/80754af9/pycparser_1594388511720/work
Pygments==2.10.0
pyOpenSSL @ file:///tmp/build/80754af9/pyopenssl_1608057966937/work
pyparsing==3.0.6
pyrsistent==0.18.0
PySocks @ file:///tmp/build/80754af9/pysocks_1605305779399/work
python-dateutil==2.8.2
pytz==2021.3
PyYAML==6.0
pyzmq==22.3.0
regex==2022.10.31
requests @ file:///tmp/build/80754af9/requests_1608241421344/work
requests-oauthlib==1.3.0
rsa==4.8
ruamel-yaml-conda @ file:///tmp/build/80754af9/ruamel_yaml_1616016699510/work
sacrebleu==2.3.1
scipy==1.9.3
Send2Trash==1.8.0
setuptools-scm==6.3.2
six @ file:///tmp/build/80754af9/six_1623709665295/work
sniffio==1.2.0
supervisor==4.2.2
sympy==1.11.1
tabulate==0.9.0
tensorboard==2.7.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.0
tensorboardX==2.6
terminado==0.12.1
testpath==0.5.0
tomli==1.2.2
torch==1.10.1+cu113
torchaudio==0.10.1+cu113
torchvision==0.11.2+cu113
tornado==6.1
tqdm @ file:///tmp/build/80754af9/tqdm_1625563689033/work
traitlets==5.1.1
triton==2.0.0
typing-extensions==4.0.0
urllib3 @ file:///tmp/build/80754af9/urllib3_1625084269274/work
wcwidth==0.2.5
webencodings==0.5.1
websocket-client==1.2.1
Werkzeug==2.0.2
widgetsnbextension==3.5.2
zipp==3.6.0

import numpy as np
import math
# 离散随机变量的KL散度和JS散度的计算方法
def KL(p, q):
    # p,q为两个list，里面存着对应的取值的概率，整个list相加为1
    if 0 in q:
        raise ValueError
    return sum(_p * math.log(_p / _q) for (_p, _q) in zip(p, q) if _p != 0)
def JS(p, q):
    M = [0.5 * (_p + _q) for (_p, _q) in zip(p, q)]
    return 0.5 * (KL(p, M) + KL(q, M))
def exp(a, b):
    a = np.array(a, dtype=np.float32)
    b = np.array(b, dtype=np.float32)
    a /= a.sum()
    b /= b.sum()
    print(a)
    print(b)
    print(KL(a, b))
    print(JS(a, b))
if __name__ == '__main__':
    # exp1
    print('exp1: Start')
    print(exp([1, 2, 3, 4, 5], [5, 4, 3, 2, 1]))
    print('exp1: End')
    # exp2
    # 把公式中的第二个分布做修改，假设这个分布中有某个值的取值非常小，就有可能增加两个分布的散度值
    print('exp2: Start')
    print(exp([1, 2, 3, 4, 5], [1e-12, 4, 3, 2, 1]))
    print(exp([1, 2, 3, 4, 5], [5, 4, 3, 2, 1e-12]))
    print('exp2: End')
    # exp3
    print('exp3: Start')
    print(exp([1e-12,2,3,4,5],[5,4,3,2,1]))
    print(exp([1,2,3,4,1e-12],[5,4,3,2,1]))
    print('exp3: End')