# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ### Beam statistics?

pred_dir = "/Users/memray/Project/keyphrase/OpenNMT-kpg/output/meng17-one2many-beam10-maxlen40/pred"

import os
import json

ckpt_list = ["kp20k-meng17-length-rnn-BS64-LR0.05-Layer1-Dim150-Emb100-Dropout0.0-Copytrue-Reusetrue-Covtrue-PEfalse-Contboth-IF1_step_90000", "kp20k-meng17-random-rnn-BS64-LR0.05-Layer1-Dim150-Emb100-Dropout0.0-Copytrue-Reusetrue-Covtrue-PEfalse-Contboth-IF1_step_90000", 
"kp20k-meng17-alphabetical-rnn-BS64-LR0.05-Layer1-Dim150-Emb100-Dropout0.0-Copytrue-Reusetrue-Covtrue-PEfalse-Contboth-IF1_step_70000", 
"kp20k-meng17-no_sort-rnn-BS64-LR0.05-Layer1-Dim150-Emb100-Dropout0.0-Copytrue-Reusetrue-Covtrue-PEfalse-Contboth-IF1_step_90000", 
"kp20k-meng17-verbatim_prepend-rnn-BS64-LR0.05-Layer1-Dim150-Emb100-Dropout0.0-Copytrue-Reusetrue-Covtrue-PEfalse-Contboth-IF1_step_65000", 
"kp20k-meng17-verbatim_append-rnn-BS64-LR0.05-Layer1-Dim150-Emb100-Dropout0.0-Copytrue-Reusetrue-Covtrue-PEfalse-Contboth-IF1_step_50000"]

datasets = ["kp20k_valid500", "duc", "inspec", "krapivin", "nus", "semeval"]

# +

for ckpt in ckpt_list:
    print(ckpt)
    
    beam_num = []
    beam_len = []
    for dataset in datasets:
        # print(dataset)
        
        pred_json_path = os.path.join(pred_dir, ckpt, dataset + '.pred')
        for jsonl in open(pred_json_path, 'r'):
            pred = json.loads(jsonl)
            beams = pred["ori_pred_sents"]
            beam_num.append(len(beams))
            beam_len.extend([len(b) for b in beams])
        
        # print("beam number: total=%d, avg=%f" % (sum(beam_num), sum(beam_num)/len(beam_num)))
        # print("beam length: total=%d, avg=%f" % (sum(beam_len), sum(beam_len)/len(beam_len)))
        # print("%d\t%f\t%d\t%f" % (sum(beam_num), sum(beam_num)/len(beam_num), sum(beam_len), sum(beam_len)/len(beam_len)))
        print(sum(beam_len))
# -

os.path.abspath(pred_dir)

os.path.exists(pred_dir)



StarcoderdataJupyterScriptsDedupFiltered