aij_test.py

from team_code.generate import get_ppl, generate_text, setup_model_and_tokenizer

def test_generate_text_get_ppl():

    model, tokenizer = setup_model_and_tokenizer()

    queries = [
        [{"type": "text", "content": "What is the capital of Great Britain?"}],
        [{"type": "text", "content": "And what is the famous place there?"}]
    ]

    resp = None
    history = None
    for cur_query_list in queries:
        resp, history = generate_text(model, tokenizer, cur_query_list, history)
        print("resp", resp)

        cur_query_tuple = ( cur_query_list, resp )
        ppl, history = get_ppl(model, tokenizer, cur_query_tuple, history)
        print("ppl", ppl)