import heapq
class Beam:
    def __init__(self):
        self.heap = list() #保存数据的位置
        self.beam_width = 3 #保存数据的总数

    def add(self,probility,complete,seq,decoder_input,decoder_hidden):
        """
        添加数据，同时判断总的数据个数，多则删除
        :param probility: 概率乘积
        :param complete: 最后一个是否为EOS
        :param seq: list，所有token的列表
        :param decoder_input: 下一次进行解码的输入，通过前一次获得
        :param decoder_hidden: 下一次进行解码的hidden，通过前一次获得
        :return:
        """
        heapq.heappush(self.heap,[probility,complete,seq,decoder_input,decoder_hidden])
        #判断数据的个数，如果大，则弹出。保证数据总个数小于等于3
        if len(self.heap)>self.beam_width:
            heapq.heappop(self.heap)

    def __iter__(self):#让该beam能够被迭代
        return iter(self.heap)