from transformers import AutoModelForCausalLM, AutoTokenizer
import torch


def load_rag_benchmark_tester_ds():

    # pull 200 question rag benchmark test dataset from LLMWare HuggingFace repo
    from datasets import load_dataset

    ds_name = "llmware/rag_instruct_benchmark_tester"

    dataset = load_dataset(ds_name)

    print("update: loading test dataset - ", dataset)

    test_set = []
    for i, samples in enumerate(dataset["train"]):
        test_set.append(samples)

        # to view test set samples
        # print("rag benchmark dataset test samples: ", i, samples)

    return test_set


def run_test(model_name, test_ds):

    device = "cuda" if torch.cuda.is_available() else "cpu"

    print("update: model will be loaded on device - ", device)

    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
    model.to(device)

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

    for i, entries in enumerate(test_ds):

        # prepare prompt packaging used in fine-tuning process
        new_prompt = "<human>: " + entries["context"] + "\n" + entries["query"] + "\n" + "<bot>:"

        inputs = tokenizer(new_prompt, return_tensors="pt")
        start_of_output = len(inputs.input_ids[0])

        #   temperature: set at 0.3 for consistency of output
        #   max_new_tokens:  set at 100 - may prematurely stop a few of the summaries

        outputs = model.generate(
            inputs.input_ids.to(device),
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True,
            temperature=0.3,
            max_new_tokens=100,
            )

        output_only = tokenizer.decode(outputs[0][start_of_output:],skip_special_tokens=True)

        # quick/optional post-processing clean-up of potential fine-tuning artifacts

        eot = output_only.find("<|endoftext|>")
        if eot > -1:
            output_only = output_only[:eot]

        bot = output_only.find("<bot>:")
        if bot > -1:
            output_only = output_only[bot+len("<bot>:"):]

        # end - post-processing

        print("\n")
        print(i, "llm_response - ", output_only)
        print(i, "gold_answer - ", entries["answer"])

    return 0


if __name__ == "__main__":

    test_ds = load_rag_benchmark_tester_ds()

    model_name = "llmware/dragon-red-pajama-7b-v0"
    output = run_test(model_name,test_ds)