# ---
# 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
# ---

import torch
import datetime
import json
from gptj_utils import GPTJ_PrefixTune

gptj = GPTJ_PrefixTune.from_pretrained(
    './model_webnlg/checkpoint-3250/',
    main_checkpoint_override="/export/data/gptj/j6b_ckpt/",
    deepspeed_config='ds_config_stage2_gptj_gen.json',
)

val_dataset_file = "/export/WebNLG/webnlg-dataset/webnlg_challenge_2017/val.json"


# +

def get_inputs(dataset_file, model):

    with open(dataset_file) as f:
        lines_dict = json.load(f)

    full_rela_lst = []
    full_src_lst = []
    full_tgt_lst = []

    for i, example in enumerate(lines_dict['entries']):
        sents = example[str(i + 1)]['lexicalisations']
        triples = example[str(i + 1)]['modifiedtripleset']

        rela_lst = []
        temp_triples = ''
        for j, tripleset in enumerate(triples):
            subj, rela, obj = tripleset['subject'], tripleset['property'], tripleset['object']
            rela_lst.append(rela)
            temp_triples += ' | '
            temp_triples += '{} : {} : {}'.format(subj, rela, obj)

        for sent in sents:
            if sent["comment"] == 'good':
                full_tgt_lst.append(sent["lex"])
                full_src_lst.append(temp_triples)
                full_rela_lst.append(rela_lst)

    edited_sents = []
    for src, tgt in zip(full_src_lst, full_tgt_lst):
        sent = ' {} {} '.format(src, model.tokenizer.bos_token)
        edited_sents.append(sent)
        
    return edited_sents


# -

sents = get_inputs(val_dataset_file, gptj)

sents[0]

eos_token_id = gptj.tokenizer.encode("<|endoftext|>")[0]
eos_token_id

idx = 0

with torch.no_grad():
    text = sents[idx]
    start = datetime.datetime.now()
    out = gptj.generate(
        1,
        text=text,
        max_length=256,
        #num_beams=2,
        do_sample=True,
        temperature=0.1,
        top_k=5,
        top_p=0.95,
        no_repeat_ngram_size=2, 
        early_stopping=True,
        num_return_sequences=3,
        use_cache=False,
        eos_token_id=eos_token_id,
        return_only_generated=True,
    )
    duration = datetime.datetime.now() - start
    for o in out:
        print(f"# =================== #\n{o[:]}\n\n\n")
    print(f"\n\nDuration = {duration.total_seconds()}")

sents[idx]


gptj.s_wte.n_prompts



StarcoderdataJupyterScriptsDedupFiltered