# Path: sockeye/beam_search.py
# class CandidateScorer(pt.nn.Module):
# 
#     def __init__(self,
#                  length_penalty_alpha: float = 1.0,
#                  length_penalty_beta: float = 0.0,
#                  brevity_penalty_weight: float = 0.0) -> None:
#         super().__init__()
#         self._lp = LengthPenalty(alpha=length_penalty_alpha, beta=length_penalty_beta)
#         self._bp = None  # type: Optional[BrevityPenalty]
#         if brevity_penalty_weight > 0.0:
#             self._bp = BrevityPenalty(weight=brevity_penalty_weight)
# 
#     def forward(self, scores, lengths, reference_lengths):
#         lp = self._lp(lengths)
#         if self._bp is not None:
#             bp = self._bp(lengths, reference_lengths)
#         else:
#             bp = 0.0
#         if isinstance(scores, (int, float)):
#             return scores / lp - bp
#         else:
#             if isinstance(lp, pt.Tensor):
#                 lp = lp.to(scores.dtype)
#             if isinstance(bp, pt.Tensor):
#                 bp = bp.to(scores.dtype)
#             return (scores.squeeze(1) / lp - bp).unsqueeze(1)
# 
#     def unnormalize(self, scores, lengths, reference_lengths):
#         bp = 0.0 if self._bp is None else self._bp(lengths, reference_lengths)
#         if isinstance(scores, (int, float)):
#             return (scores + bp) * self._lp(lengths)
#         else:
#             return ((scores.squeeze(1) + bp) * self._lp(lengths)).unsqueeze(1)


import torch as pt
import sockeye.scoring
from sockeye.beam_search import CandidateScorer
# Copyright 2017--2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You may not
# use this file except in compliance with the License. A copy of the License
# is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is distributed on
# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.




def test_batch_scorer():
    # TODO: make this a useful test
    batch = 2
    seq = 4
    nh = 6
    logits = pt.ones(batch, seq, nh)
    label = pt.ones(batch, seq).to(pt.long)
    length_ratio = pt.ones(batch, )
    source_length = pt.randint(0, seq, (batch,)).to(pt.float32)
    target_length = source_length

# Path: sockeye/beam_search.py
# class CandidateScorer(pt.nn.Module):
# 
#     def __init__(self,
#                  length_penalty_alpha: float = 1.0,
#                  length_penalty_beta: float = 0.0,
#                  brevity_penalty_weight: float = 0.0) -> None:
#         super().__init__()
#         self._lp = LengthPenalty(alpha=length_penalty_alpha, beta=length_penalty_beta)
#         self._bp = None  # type: Optional[BrevityPenalty]
#         if brevity_penalty_weight > 0.0:
#             self._bp = BrevityPenalty(weight=brevity_penalty_weight)
# 
#     def forward(self, scores, lengths, reference_lengths):
#         lp = self._lp(lengths)
#         if self._bp is not None:
#             bp = self._bp(lengths, reference_lengths)
#         else:
#             bp = 0.0
#         if isinstance(scores, (int, float)):
#             return scores / lp - bp
#         else:
#             if isinstance(lp, pt.Tensor):
#                 lp = lp.to(scores.dtype)
#             if isinstance(bp, pt.Tensor):
#                 bp = bp.to(scores.dtype)
#             return (scores.squeeze(1) / lp - bp).unsqueeze(1)
# 
#     def unnormalize(self, scores, lengths, reference_lengths):
#         bp = 0.0 if self._bp is None else self._bp(lengths, reference_lengths)
#         if isinstance(scores, (int, float)):
#             return (scores + bp) * self._lp(lengths)
#         else:
#             return ((scores.squeeze(1) + bp) * self._lp(lengths)).unsqueeze(1)


# Path: test/unit/test_scoring.py
import torch as pt
import sockeye.scoring
from sockeye.beam_search import CandidateScorer

# Copyright 2017--2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You may not
# use this file except in compliance with the License. A copy of the License
# is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is distributed on
# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.




def test_batch_scorer():
    # TODO: make this a useful test
    batch = 2
    seq = 4
    nh = 6
    logits = pt.ones(batch, seq, nh)
    label = pt.ones(batch, seq).to(pt.long)
    length_ratio = pt.ones(batch, )
    source_length = pt.randint(0, seq, (batch,)).to(pt.float32)
    target_length = source_length

    b = sockeye.scoring.BatchScorer(scorer=CandidateScorer(1.0, 0.0, 0.0),