(max_wn, s_wv, beam_size)
return pr_wvi_beam, prob_wvi_beam
s_wv: [B, 4, mL, 2]
    - predict best st-idx & ed-idx


    output:
    pr_wvi_beam = [B, max_wn, n_pairs, 2]. 2 means [st, ed].
    prob_wvi_beam = [B, max_wn, n_pairs]
s_wv: [B, 4, mL, 2]
    - predict best st-idx & ed-idx
s_wv : [ B 4 mL 2 ] - predict best st - idx & ed - idx
def pred_wvi_se_beam(max_wn, s_wv, beam_size):
    """
    s_wv: [B, 4, mL, 2]
    - predict best st-idx & ed-idx


    output:
    pr_wvi_beam = [B, max_wn, n_pairs, 2]. 2 means [st, ed].
    prob_wvi_beam = [B, max_wn, n_pairs]
    """
    bS = s_wv.shape[0]

    s_wv_st, s_wv_ed = s_wv.split(1, dim=3)  # [B, 4, mL, 2] -> [B, 4, mL, 1], [B, 4, mL, 1]

    s_wv_st = s_wv_st.squeeze(3) # [B, 4, mL, 1] -> [B, 4, mL]
    s_wv_ed = s_wv_ed.squeeze(3)

    prob_wv_st = F.softmax(s_wv_st, dim=-1).detach().to('cpu').numpy()
    prob_wv_ed = F.softmax(s_wv_ed, dim=-1).detach().to('cpu').numpy()

    k_logit = int(ceil(sqrt(beam_size)))
    n_pairs = k_logit**2
    assert n_pairs >= beam_size
    values_st, idxs_st = s_wv_st.topk(k_logit) # [B, 4, mL] -> [B, 4, k_logit]
    values_ed, idxs_ed = s_wv_ed.topk(k_logit) # [B, 4, mL] -> [B, 4, k_logit]

    # idxs = [B, k_logit, 2]
    # Generate all possible combination of st, ed indices & prob
    pr_wvi_beam = [] # [B, max_wn, k_logit**2 [st, ed] paris]
    prob_wvi_beam = zeros([bS, max_wn, n_pairs])
    for b in range(bS):
        pr_wvi_beam1 = []

        idxs_st1 = idxs_st[b]
        idxs_ed1 = idxs_ed[b]
        for i_wn in range(max_wn):
            idxs_st11 = idxs_st1[i_wn]
            idxs_ed11 = idxs_ed1[i_wn]

            pr_wvi_beam11 = []
            pair_idx = -1
            for i_k in range(k_logit):
                for j_k in range(k_logit):
                    pair_idx += 1
                    st = idxs_st11[i_k].item()
                    ed = idxs_ed11[j_k].item()
                    pr_wvi_beam11.append([st, ed])

                    p1 = prob_wv_st[b, i_wn, st]
                    p2 = prob_wv_ed[b, i_wn, ed]
                    prob_wvi_beam[b, i_wn, pair_idx] = p1*p2
            pr_wvi_beam1.append(pr_wvi_beam11)
        pr_wvi_beam.append(pr_wvi_beam1)


    # prob

    return pr_wvi_beam, prob_wvi_beam
def pred_wvi_se_beam ( max_wn , s_wv , beam_size ) : bS = s_wv . shape [ 0 ] s_wv_st , s_wv_ed = s_wv . split ( 1 , dim = 3 ) # [B, 4, mL, 2] -> [B, 4, mL, 1], [B, 4, mL, 1] s_wv_st = s_wv_st . squeeze ( 3 ) # [B, 4, mL, 1] -> [B, 4, mL] s_wv_ed = s_wv_ed . squeeze ( 3 ) prob_wv_st = F . softmax ( s_wv_st , dim = - 1 ) . detach ( ) . to ( 'cpu' ) . numpy ( ) prob_wv_ed = F . softmax ( s_wv_ed , dim = - 1 ) . detach ( ) . to ( 'cpu' ) . numpy ( ) k_logit = int ( ceil ( sqrt ( beam_size ) ) ) n_pairs = k_logit ** 2 assert n_pairs >= beam_size values_st , idxs_st = s_wv_st . topk ( k_logit ) # [B, 4, mL] -> [B, 4, k_logit] values_ed , idxs_ed = s_wv_ed . topk ( k_logit ) # [B, 4, mL] -> [B, 4, k_logit] # idxs = [B, k_logit, 2] # Generate all possible combination of st, ed indices & prob pr_wvi_beam = [ ] # [B, max_wn, k_logit**2 [st, ed] paris] prob_wvi_beam = zeros ( [ bS , max_wn , n_pairs ] ) for b in range ( bS ) : pr_wvi_beam1 = [ ] idxs_st1 = idxs_st [ b ] idxs_ed1 = idxs_ed [ b ] for i_wn in range ( max_wn ) : idxs_st11 = idxs_st1 [ i_wn ] idxs_ed11 = idxs_ed1 [ i_wn ] pr_wvi_beam11 = [ ] pair_idx = - 1 for i_k in range ( k_logit ) : for j_k in range ( k_logit ) : pair_idx += 1 st = idxs_st11 [ i_k ] . item ( ) ed = idxs_ed11 [ j_k ] . item ( ) pr_wvi_beam11 . append ( [ st , ed ] ) p1 = prob_wv_st [ b , i_wn , st ] p2 = prob_wv_ed [ b , i_wn , ed ] prob_wvi_beam [ b , i_wn , pair_idx ] = p1 * p2 pr_wvi_beam1 . append ( pr_wvi_beam11 ) pr_wvi_beam . append ( pr_wvi_beam1 ) # prob return pr_wvi_beam , prob_wvi_beam