Spaces:
Sleeping
Sleeping
import torch | |
import nltk | |
from nltk.translate.bleu_score import SmoothingFunction | |
from tqdm import tqdm | |
def calculate_perplexity(model, tokens, prompt_len, bsz=1, marker=False): | |
""" | |
Calculate perplexity of given tokens using provided model, ignoring padding tokens. | |
Args: | |
model: Llama model | |
tokens (List[List[int]] or torch.Tensor): Input tokens (n_prompt * n_draft, seqlen) | |
prompt_len (int): Prefix length | |
bsz (int): Batch size | |
marker (bool): Whether to show progress bar | |
Returns: | |
Perplexity across all generations (n_prompt * n_drafts) | |
""" | |
it = range(0, len(tokens), bsz) | |
if marker: | |
it = tqdm(it) | |
start = 0 | |
ppl = torch.zeros(len(tokens)) | |
for start in it: | |
end = start + bsz | |
data = tokens[start : end] | |
if not isinstance(data, list): | |
data = data.tolist() | |
# Remove any padding tokens (-1) in generations | |
for d_idx in range(len(data)): | |
cur = data[d_idx] | |
if -1 in cur: | |
data[d_idx] = cur[:cur.index(-1)] | |
# Calculate cross entropy loss on tokens | |
ce_loss = model.generate(data, max_gen_len=0, temperature=-1, top_p=-1, grade=True) | |
# Cut off everything past `prompt_len` | |
ce_loss = ce_loss[:, prompt_len-1:] # Subtract 1 because the first token (start token) is removed | |
# Calculate perplexity | |
lengths = (ce_loss != 0).sum(dim=-1) | |
mean = ce_loss.sum(dim=-1) / lengths | |
ppl[start : end] = torch.exp(-1 * mean) | |
return ppl | |
def calculate_diversity(generations, k=4): | |
""" | |
Calculate diversity of generations using SELF-BLEU. | |
Args: | |
generations (List[List[List[int]]]): Tokenized input | |
k (int, Optional): Number of n-grams to use for bleu | |
Returns: | |
Average diversity across all generations (float) | |
""" | |
nltk.download('punkt') # Can be deleted once downloaded | |
smooth = SmoothingFunction() | |
bleus = [] | |
for drafts in generations: | |
tokenized_drafts = [] | |
# Stringify tokens | |
for d in drafts: | |
if -1 in d: | |
d = d[:d.index(-1)] | |
tokenized_drafts.append([str(n) for n in d]) | |
# Calculate SELF-BLEU | |
minlength = min([len(g) for g in tokenized_drafts]) | |
minlength = min(minlength, k) | |
weights = tuple((1. / minlength for _ in range(minlength))) | |
for i in range(len(drafts)): | |
# Create source and reference (all other drafts) | |
src = tokenized_drafts[i] | |
ref = tokenized_drafts[:i] + tokenized_drafts[i+1:] | |
tmp = nltk.translate.bleu_score.sentence_bleu(references=ref, | |
hypothesis=src, | |
weights=weights, | |
smoothing_function=smooth.method1) | |
bleus.append(tmp) | |
bleus = torch.Tensor(bleus) | |
return torch.mean(bleus) | |
def calculate_ngram_repetition(sequences): | |
""" | |
Calculate uniqueness scores of `sequences`. | |
Args: | |
sequences (List[List[int]]): Generated sequences | |
Returns: | |
(unigram_uniqueness, bigram_uniqueness, trigram_uniqueness) | |
""" | |
u_total = 0 | |
b_total = 0 | |
t_total = 0 | |
# Iterate through all sequences indiscriminately | |
for gen in sequences: | |
if -1 in gen: | |
gen = gen[:gen.index(-1)] | |
unigrams, bigrams, trigrams = [], [], [] | |
o = [str(i) for i in gen] | |
# Create lists of n-grams for the generation | |
for i in range(len(o)): | |
unigrams.append(o[i]) | |
for i in range(len(o) - 1): | |
bigrams.append(o[i] + '_' + o[i + 1]) | |
for i in range(len(o) - 2): | |
trigrams.append(o[i] + '_' + o[i + 1] + '_' + o[i + 2]) | |
# Calculate uniqueness of the generation | |
u, b, t = len(set(unigrams)) / len(unigrams), len(set(bigrams)) / len(bigrams), len(set(trigrams)) / len(trigrams) | |
u_total += u | |
b_total += b | |
t_total += t | |
return u_total / len(sequences), b_total / len(sequences), t_total / len(sequences) | |