from sklearn.model_selection import train_test_split

from headliner.training_runner import read_data_json
from headliner.model.summarizer_bert import SummarizerBert

if __name__ == '__main__':

    train = read_data_json('/Users/cschaefe/datasets/welt_dedup.json', 400)[:1000]
    train_data, val_data = train_test_split(train, test_size=100, shuffle=True, random_state=32)
    path_to_model = '/Users/cschaefe/transformer_check_7'
    summarizer = SummarizerBert.load(path_to_model)
    summarizer.max_prediction_len = 30
    for val_index in range(100):
        text = val_data[val_index][0]
        target = val_data[val_index][1]
        print('(target) {}'.format(target))
        pred = summarizer.predict_beam_search(text, beam_width=15)
        for p in pred:
            print('{} {}'.format(p[0], p[1]))
