%%time
pipe = pipeline("text-generation", model="gpt2", max_new_tokens=128, do_sample=True)
prompt = """ Universal Basic Income, also known as UBI, is """
res = pipe(prompt)

p80(res[0]['generated_text'])

%%time

def load_model_tokenizer(model_name_or_path):
    config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path, config=config, trust_remote_code=True, cache_dir = f"cache/{model_name_or_path.split('/')[-1]}",
        # load_in_4bit=True,
        # load_in_8bit=True,
        device_map="auto",
        offload_folder=f"offloads/{model_name_or_path.split('/')[-1]}"
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path,
        padding_side="left",
        cache_dir=f"cache/{model_name_or_path.split('/')[-1]}"
    )
    tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer

# LLM
# model_name = "Writer/palmyra-small" # https://huggingface.co/Writer/palmyra-small
# model_name = "mistralai/Mistral-7B-v0.1" # https://huggingface.co/mistralai/Mistral-7B-v0.1
# model_name = "HuggingFaceH4/zephyr-7b-alpha" # https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha
model_name="microsoft/phi-1_5" # https://huggingface.co/microsoft/phi-1_5

model, tokenizer = load_model_tokenizer(model_name)

def generate_response(prompt, tokenizer, model):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    generate_ids = model.generate(inputs.input_ids, max_new_tokens=512, do_sample=True)
    response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
    response = response.split(prompt)[-1]
    return response

%%time
res = generate_response(prompt, tokenizer, model)
p80(prompt + res)

def load_model_tokenizer(model_name_or_path):
    config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path, config=config, trust_remote_code=True, cache_dir = f"cache/{model_name_or_path.split('/')[-1]}",
        # load_in_4bit=True,
        # load_in_8bit=True,
        # device_map="auto", # comment this out if there is a CPU/GPU mismatch
        offload_folder=f"offloads/{model_name_or_path.split('/')[-1]}"
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path,
        padding_side="left",
        cache_dir=f"cache/{model_name_or_path.split('/')[-1]}"
    )
    tokenizer.pad_token = tokenizer.eos_token
    return model, tokenizer

# LLM
model_name = "Writer/palmyra-small" # https://huggingface.co/Writer/palmyra-small
# model_name = "mistralai/Mistral-7B-v0.1" # https://huggingface.co/mistralai/Mistral-7B-v0.1
# model_name = "HuggingFaceH4/zephyr-7b-alpha" # https://huggingface.co/HuggingFaceH4/zephyr-7b-alpha
# model_name="microsoft/phi-1_5" # https://huggingface.co/microsoft/phi-1_5

model, tokenizer = load_model_tokenizer(model_name)

import torch
from tqdm import tqdm

def get_ppl(text, model, tokenizer, stride=512):
    encodings = tokenizer(text, return_tensors="pt")
    max_length = stride #model.config.n_positions
    seq_len = encodings.input_ids.size(1)
    nlls = []
    prev_end_loc = 0
    for begin_loc in range(0, seq_len, stride):
        end_loc = min(begin_loc + max_length, seq_len)
        trg_len = end_loc - prev_end_loc  # may be different from stride on last loop
        input_ids = encodings.input_ids[:, begin_loc:end_loc]#.to(device)
        target_ids = input_ids.clone()
        target_ids[:, :-trg_len] = -100

        with torch.no_grad():
            outputs = model(input_ids, labels=target_ids)

            # loss is calculated using CrossEntropyLoss which averages over valid labels
            # N.B. the model only calculates loss over trg_len - 1 labels, because it internally shifts the labels
            # to the left by 1.
            neg_log_likelihood = outputs.loss

        nlls.append(neg_log_likelihood)

        prev_end_loc = end_loc
        if end_loc == seq_len:
            break

    ppl = torch.exp(torch.stack(nlls).mean())
    return ppl.numpy().tolist()

## EXAMPLE 1

# model_name = "gpt2"
# model, tokenizer = load_model_tokenizer(model_name)
get_ppl(""" I love PyTorch """, model, tokenizer)

## EXAMPLE 2
txt = """ 'Alexis Sanchez had offers from several clubs when he left Barcelona last summer, including Arsenal and Liverpool. He chose to join Arsenal in a decision that delighted boss Arsene Wenger. Sanchez has made an instant impact in English football and has scored 19 goals for the Gunners so far this season. Arsenal face Liverpool in the Premier League on Saturday as the two sides compete for a top-four finish.', 'Luis Suarez the famous chilean striker has departed Barcelona FC and is currently being courted by teams in the Barclays Premier League. Manchester United is the most likely landing spot for the famous Chilean striker who has scored hundreds of goals throughout his famous career. Negotiations between interested parties are ongoing and a decision is expected to be heard in the coming weeks.', 'An impressive english football player is being scouted out by several teams that hope to have him on their team. These bidding wars are intense and the player still has yet to make a decision.', 'Wenger is happy to have Alexis Sanchez. Sanchez was previously dismissed by Barcelona. He has been influential with 19 goals this season.', 'Arsenal Manager Arsene Wenger elated that Alexis Sanchez will side with his team for Liverpool on Saturday.', '26 year old Alexis Sanchez will be lining up for Arsenal this Saturday, and manager Arsene Wenger is certainly happy about it. When Barcelona allowed Sanchez to leave, he could have chosen any team. Luckily, he chose Arsenal and has already scored nineteen goals this season.', 'Alexis Sanchez chose Emirates over Anfield. Alexis Sanchez impressed immediately with his stellar play. Wenger has been rumored to be linked with Liverpool winger Raheem Sterling.', 'Alexis Sanchez picked Emirates over Anfield. Alexis Sanchez has made an instant impact with 19 goals. The speculation that is being linked with Wenger is Raheem Sterling.', "Chile forward Alexis Sanchez recently chose to play for Arsenal at the Emirates Statium over Liverpool at Anfield. Sanchez came into the league strong, scoring 19 goals this season. Arsenal manager Arsene Wenger is trying to downplay speculation that links him to Liverpool's winger Raheem Sterling.", 'Alexus Sabcgez chose the play at Emerates Stadium over Anfield in the premier league. Sanchez has mad a great impact with the teal scoring 19 goals this year for the North London Club. Wenger has been impressed by the way Rodger has developed his team but has down played speculation linking him to Raheem Sterling, the Liverpool winger.', "Alexis Sanchez picked Arsenal and Emirates Stadium to play for instead of Liverpool and Anfield, and his coach Arsene Wenger is sure happy about it. Sanchez has already scored nineteen goals for his team this year. Wenger also downplayed an interest in Raheem Sterling, who is a forward for Liverpool. Sterling rejected his team's latest contract offer." """

print("INPUT:")
p80(txt)

print("\nPPL metric =", get_ppl(txt, model, tokenizer))