import torch
import torch.nn as nn
import torch.nn.functional as F

class BeamSearchDecoder(nn.Module):
    def __init__(self, beam_width):
        super(BeamSearchDecoder, self).__init__()
        self.beam_width = beam_width

    def forward(self, inputs):
        batch_size, input_length, hidden_size = inputs.size()

        # 初始化束搜索队列
        beam_queue = [(torch.zeros(batch_size, hidden_size), [torch.tensor(0)]*input_length)]

        for i in range(input_length):
            new_beam_queue = []
            for score, seq in beam_queue:
                output = self.generate_output(inputs, seq)
                new_beam_queue.extend(self.expand_beam(score, seq, output))
            new_beam_queue.sort(key=lambda x: x[0], reverse=True)
            beam_queue = new_beam_queue[:self.beam_width]

        return beam_queue[0][1]  # 输出最优的序列

    def generate_output(self, inputs, seq):
        # 根据输入序列和当前生成序列生成输出
        # 这里使用了简单的线性映射来生成输出
        output = torch.matmul(seq[-1].unsqueeze(1), inputs.transpose(1, 2)).squeeze(1)
        return output

    def expand_beam(self, score, seq, output):
        # 根据当前得分、当前序列和输出来扩展束搜索队列
        candidates = [(score + F.log_softmax(output[i]), seq + [torch.tensor(i)]) for i in range(output.size(0))]
        return candidates

import torch

inputs = torch.tensor([[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]],
                       [[0.9, 1.0, 1.1, 1.2], [1.3, 1.4, 1.5, 1.6]]])  # 批次大小为2

beam_search_decoder = BeamSearchDecoder(3)
outputs = beam_search_decoder(inputs)

print(outputs)

[[tensor([1, 2, 2, 3]), tensor([0, 0, 0, 1]), tensor([2, 1, 2, 0])],
 [tensor([0, 3, 3, 2]), tensor([1, 3, 3, 2]), tensor([1, 2, 3, 3])]]