Upload folder using huggingface_hub

c40457e verified 1 day ago

9.84 kB

import torch
import onnxruntime

PAD_ID, MASK_ID, USER_ID, BOT_ID, EOS_ID = 0, 1, 2, 3, 4

def ngram(seq, n):
    nxt, prv = dict(), dict()

    for i in range(0, len(seq) - n + 1):
        phrase_to_right = seq[i:i + n - 1]
        phrase_to_left = seq[i + 1:i + n]
        to_right_target = seq[i + n - 1]
        to_left_target = seq[i]
        if not(PAD_ID in phrase_to_right or PAD_ID == to_right_target):
            nxt.setdefault(tuple(phrase_to_right), set()).add(to_right_target)
        if not(PAD_ID in phrase_to_left or PAD_ID == to_left_target):
            prv.setdefault(tuple(phrase_to_left), set()).add(to_left_target)
    
    return nxt, prv

def penalize_ngram(ngram_end, ngram_start, ngram_len, ans_candidate_final, candidate_eos_idx, ans_candidate_logits, ngram_penalty, ANSWER_LEN):

    for i in range(candidate_eos_idx - ngram_len + 1):
        substring_for_ngram_to_end = ans_candidate_final[0, i:i + ngram_len - 1].tolist()
        substring_for_ngram_to_start = ans_candidate_final[0, i + 1:i + ngram_len].tolist()

        end_cands_to_penalize = ngram_end.get(tuple(substring_for_ngram_to_end), None)
        start_cands_to_penalize = ngram_start.get(tuple(substring_for_ngram_to_start), None)

        if end_cands_to_penalize is not None and MASK_ID not in substring_for_ngram_to_end:
            for word in end_cands_to_penalize:
                ans_candidate_logits[:, i + ngram_len - 1, word] -= ngram_penalty

        if start_cands_to_penalize is not None and MASK_ID not in substring_for_ngram_to_start:
            for word in start_cands_to_penalize:
                ans_candidate_logits[:, i, word] -= ngram_penalty

def gstar(device, model, margin_for_gstar_replace, start_block_idx, end_block_idx, KV_cache, KV_mask,
          ans_candidate_final, context_ids, rounds, is_global=False):
    block_size = end_block_idx - start_block_idx
    indexer = torch.arange(start_block_idx, end_block_idx, device=device)  # [k] — позиции в ответе
    rows = torch.arange(block_size, device=device)                         # [k] — номера копий

    # last user phrase between last <user> and last <bot>
    context = context_ids[0].tolist()
    last_user = len(context) - 1 - context[::-1].index(USER_ID) if USER_ID in context else -1
    last_bot = len(context) - 1 - context[::-1].index(BOT_ID)
    last_user_phrase = set(context[last_user + 1:last_bot])

    def leave_one_out_probs():
        batch = ans_candidate_final.repeat(block_size, 1)            # [k, 64] — k answer copies
        batch[rows, indexer] = MASK_ID                               # pos indexer[r] masked in r_i copy
        temp_logits = model.decode(batch, KV_cache.expand(block_size, -1, -1),
                                   KV_mask.expand(block_size, -1, -1)).float()   # [k, 64, V]
        temp_logits = temp_logits[rows, indexer]                     # [k, V] only masked pos
        temp_logits[:, (MASK_ID, PAD_ID, EOS_ID)] = -float('inf')
        return torch.softmax(temp_logits, dim=-1)

    for step in range(rounds):
        temp_probs = leave_one_out_probs()
        cur_token = ans_candidate_final[0, indexer]                  # [k]
        p_cur = temp_probs[rows, cur_token]                          # [k]
        p_best, best_token = temp_probs.max(dim=-1)                  # [k], [k]

        token_to_replace, pr, diff = None, None, -float('inf')
        for pretender in range(block_size):
            is_bad = p_cur[pretender] < margin_for_gstar_replace * p_best[pretender]
            is_other = best_token[pretender] != cur_token[pretender]
            is_free = cur_token[pretender].item() not in last_user_phrase
            gap = (p_best[pretender] - p_cur[pretender]).item()
            if is_bad and is_other and is_free and gap > diff:
                token_to_replace, pr, diff = best_token[pretender], pretender, gap

        if token_to_replace is None:
            break
        ans_candidate_final[0, indexer[pr]] = token_to_replace

    if is_global:
        temp_probs = leave_one_out_probs()
        p_cur = temp_probs[rows, ans_candidate_final[0, indexer]]
        return p_cur.clamp_min(1e-9).log().mean().item()

@torch.no_grad()
def generate(model, 
             context_ids, 
             mean_length_scale, 
             min_len, 
             n_candidates, 
             block_len, 
             repeat_penalty, 
             ngram_ans_len, 
             ngram_user_len, 
             ngram_penalty, 
             temperature,
             top_k,
             margin_for_gstar_replace,
             ANSWER_LEN,
             block_rounds,
             global_rounds,
             greedy,
             device
            ):
    context_ids = context_ids[None, :] # [256] -> [1, 256]
    KV_cache, KV_mask = model.encode(context_ids) # [1, 256, 384], [1, 1, 384]
    #answer_logits = torch.ones(size=(1, 64), dtype=torch.long) [1, 64]
    answer_probs = torch.nn.functional.softmax(
        model.decode(torch.full((1, ANSWER_LEN), MASK_ID, dtype=torch.long, device=device), KV_cache, KV_mask).float(), dim=2) # [1, 64, 12288]
    
    Len_mean = ((1 - answer_probs[..., EOS_ID] - answer_probs[..., PAD_ID]).sum(dim=1) * mean_length_scale).item()
    length_scale_candidates = torch.linspace(torch.tensor(0.8), torch.tensor(1.2), n_candidates) if n_candidates > 1 else torch.tensor([1.0])
    Len_candidates = ((Len_mean * length_scale_candidates).round()).clamp(min_len, 63).long() # [n_candidates]

    ngram_end_user, ngram_start_user = ngram(context_ids[0].tolist(), ngram_user_len)

    best_candidate, best_candidate_score = None, -float('inf')
    for idx in range(n_candidates):
        ans_candidate_final = torch.full((1, ANSWER_LEN), MASK_ID, dtype=torch.long, device=device)
        ans_candidate_final[:, Len_candidates[idx]] = EOS_ID
        ans_candidate_final[:, Len_candidates[idx] + 1:] = PAD_ID

        for start_block_idx in range(0, Len_candidates[idx].item(), block_len):
            end_block_idx = min(start_block_idx + block_len, Len_candidates[idx].item())

            block_region = torch.zeros(size=(64,), device=device).bool()
            block_region[start_block_idx:end_block_idx] = True

            while True:
                masks_still_to_open = (ans_candidate_final[0] == MASK_ID) & block_region
                if not masks_still_to_open.any().item(): break

                ans_candidate_logits = model.decode(ans_candidate_final, KV_cache, KV_mask).float()

                #penalty wrong end
                ans_candidate_logits[..., (MASK_ID, PAD_ID, EOS_ID)] = -float('inf')

                #penalty repeat
                used = ans_candidate_final[0][ans_candidate_final[0] != MASK_ID].unique()
                lu = ans_candidate_logits[0, :, used]
                ans_candidate_logits[0, :, used] = torch.where(lu > 0, lu / repeat_penalty, lu * repeat_penalty)

                #penalty ngram
                ngrams_ans_end, ngrams_ans_start = ngram(ans_candidate_final[0].tolist(), ngram_ans_len)
                penalize_ngram(ngrams_ans_end, ngrams_ans_start, ngram_ans_len, ans_candidate_final, Len_candidates[idx].item(), 
                               ans_candidate_logits, ngram_penalty, ANSWER_LEN)
                penalize_ngram(ngram_end_user, ngram_start_user, ngram_user_len, ans_candidate_final, Len_candidates[idx].item(), 
                               ans_candidate_logits, ngram_penalty, ANSWER_LEN)

                # temperature
                ans_candidate_logits[0] /= temperature

                #top k
                min_logit_per_row = ans_candidate_logits.topk(top_k, dim=2).values[:, :, -1:] # [64, 1]
                ans_candidate_logits = ans_candidate_logits.masked_fill(ans_candidate_logits < min_logit_per_row, -float("inf"))

                # probs and word choosing
                ans_candidate_probs = torch.softmax(ans_candidate_logits, 2).squeeze(0) # [1, 64, 12288] -> [64, 12288]

                if greedy:
                    tokens_choice = torch.argmax(ans_candidate_probs, dim=1, keepdim=True)
                else:
                    tokens_choice = torch.multinomial(ans_candidate_probs, 1) # [64, 1]
                confidence = ans_candidate_probs.gather(-1, tokens_choice).squeeze(1) # [64, 1] -> [64]

                # open only one mask per step
                i = confidence.masked_fill(~masks_still_to_open, -1).argmax().item()
                ans_candidate_final[0, i] = tokens_choice[i, 0]

            # gstar
            gstar(device, model, margin_for_gstar_replace, start_block_idx, end_block_idx, KV_cache, KV_mask, ans_candidate_final, context_ids, 
                  rounds=block_rounds)
        this_candidate_score = gstar(device, model, margin_for_gstar_replace, 0, Len_candidates[idx].item(), KV_cache, KV_mask, ans_candidate_final, 
                                     context_ids, rounds=global_rounds, is_global=True)
        if this_candidate_score > best_candidate_score:
            best_candidate_score = this_candidate_score
            best_candidate = ans_candidate_final
            
    return best_candidate.squeeze(0)

class OrtBackend():

    def __init__(self, enc_path, dec_path):
        self.encoder = onnxruntime.InferenceSession(enc_path, providers=['CPUExecutionProvider'])
        self.decoder = onnxruntime.InferenceSession(dec_path, providers=['CPUExecutionProvider'])

    def encode(self, context_ids):
        # context_ids [1, 256] torch.tensor
        memory, mask = self.encoder.run(None, {'context_ids': context_ids.numpy()})
        return torch.from_numpy(memory), torch.from_numpy(mask)

    def decode(self, answer_ids, memory, mask):
        # 
        return torch.from_numpy(self.decoder.run(None, {'memory': memory[:1].numpy(), 'mask': mask[:1].numpy(), 'answer_ids': answer_ids.numpy()})[0])


	import torch
	import onnxruntime

	PAD_ID, MASK_ID, USER_ID, BOT_ID, EOS_ID = 0, 1, 2, 3, 4

	def ngram(seq, n):
	nxt, prv = dict(), dict()

	for i in range(0, len(seq) - n + 1):
	phrase_to_right = seq[i:i + n - 1]
	phrase_to_left = seq[i + 1:i + n]
	to_right_target = seq[i + n - 1]
	to_left_target = seq[i]
	if not(PAD_ID in phrase_to_right or PAD_ID == to_right_target):
	nxt.setdefault(tuple(phrase_to_right), set()).add(to_right_target)
	if not(PAD_ID in phrase_to_left or PAD_ID == to_left_target):
	prv.setdefault(tuple(phrase_to_left), set()).add(to_left_target)

	return nxt, prv

	def penalize_ngram(ngram_end, ngram_start, ngram_len, ans_candidate_final, candidate_eos_idx, ans_candidate_logits, ngram_penalty, ANSWER_LEN):

	for i in range(candidate_eos_idx - ngram_len + 1):
	substring_for_ngram_to_end = ans_candidate_final[0, i:i + ngram_len - 1].tolist()
	substring_for_ngram_to_start = ans_candidate_final[0, i + 1:i + ngram_len].tolist()

	end_cands_to_penalize = ngram_end.get(tuple(substring_for_ngram_to_end), None)
	start_cands_to_penalize = ngram_start.get(tuple(substring_for_ngram_to_start), None)

	if end_cands_to_penalize is not None and MASK_ID not in substring_for_ngram_to_end:
	for word in end_cands_to_penalize:
	ans_candidate_logits[:, i + ngram_len - 1, word] -= ngram_penalty

	if start_cands_to_penalize is not None and MASK_ID not in substring_for_ngram_to_start:
	for word in start_cands_to_penalize:
	ans_candidate_logits[:, i, word] -= ngram_penalty

	def gstar(device, model, margin_for_gstar_replace, start_block_idx, end_block_idx, KV_cache, KV_mask,
	ans_candidate_final, context_ids, rounds, is_global=False):
	block_size = end_block_idx - start_block_idx
	indexer = torch.arange(start_block_idx, end_block_idx, device=device) # [k] — позиции в ответе
	rows = torch.arange(block_size, device=device) # [k] — номера копий

	# last user phrase between last <user> and last <bot>
	context = context_ids[0].tolist()
	last_user = len(context) - 1 - context[::-1].index(USER_ID) if USER_ID in context else -1
	last_bot = len(context) - 1 - context[::-1].index(BOT_ID)
	last_user_phrase = set(context[last_user + 1:last_bot])

	def leave_one_out_probs():
	batch = ans_candidate_final.repeat(block_size, 1) # [k, 64] — k answer copies
	batch[rows, indexer] = MASK_ID # pos indexer[r] masked in r_i copy
	temp_logits = model.decode(batch, KV_cache.expand(block_size, -1, -1),
	KV_mask.expand(block_size, -1, -1)).float() # [k, 64, V]
	temp_logits = temp_logits[rows, indexer] # [k, V] only masked pos
	temp_logits[:, (MASK_ID, PAD_ID, EOS_ID)] = -float('inf')
	return torch.softmax(temp_logits, dim=-1)

	for step in range(rounds):
	temp_probs = leave_one_out_probs()
	cur_token = ans_candidate_final[0, indexer] # [k]
	p_cur = temp_probs[rows, cur_token] # [k]
	p_best, best_token = temp_probs.max(dim=-1) # [k], [k]

	token_to_replace, pr, diff = None, None, -float('inf')
	for pretender in range(block_size):
	is_bad = p_cur[pretender] < margin_for_gstar_replace * p_best[pretender]
	is_other = best_token[pretender] != cur_token[pretender]
	is_free = cur_token[pretender].item() not in last_user_phrase
	gap = (p_best[pretender] - p_cur[pretender]).item()
	if is_bad and is_other and is_free and gap > diff:
	token_to_replace, pr, diff = best_token[pretender], pretender, gap

	if token_to_replace is None:
	break
	ans_candidate_final[0, indexer[pr]] = token_to_replace

	if is_global:
	temp_probs = leave_one_out_probs()
	p_cur = temp_probs[rows, ans_candidate_final[0, indexer]]
	return p_cur.clamp_min(1e-9).log().mean().item()

	@torch.no_grad()
	def generate(model,
	context_ids,
	mean_length_scale,
	min_len,
	n_candidates,
	block_len,
	repeat_penalty,
	ngram_ans_len,
	ngram_user_len,
	ngram_penalty,
	temperature,
	top_k,
	margin_for_gstar_replace,
	ANSWER_LEN,
	block_rounds,
	global_rounds,
	greedy,
	device
	):
	context_ids = context_ids[None, :] # [256] -> [1, 256]
	KV_cache, KV_mask = model.encode(context_ids) # [1, 256, 384], [1, 1, 384]
	#answer_logits = torch.ones(size=(1, 64), dtype=torch.long) [1, 64]
	answer_probs = torch.nn.functional.softmax(
	model.decode(torch.full((1, ANSWER_LEN), MASK_ID, dtype=torch.long, device=device), KV_cache, KV_mask).float(), dim=2) # [1, 64, 12288]

	Len_mean = ((1 - answer_probs[..., EOS_ID] - answer_probs[..., PAD_ID]).sum(dim=1) * mean_length_scale).item()
	length_scale_candidates = torch.linspace(torch.tensor(0.8), torch.tensor(1.2), n_candidates) if n_candidates > 1 else torch.tensor([1.0])
	Len_candidates = ((Len_mean * length_scale_candidates).round()).clamp(min_len, 63).long() # [n_candidates]

	ngram_end_user, ngram_start_user = ngram(context_ids[0].tolist(), ngram_user_len)

	best_candidate, best_candidate_score = None, -float('inf')
	for idx in range(n_candidates):
	ans_candidate_final = torch.full((1, ANSWER_LEN), MASK_ID, dtype=torch.long, device=device)
	ans_candidate_final[:, Len_candidates[idx]] = EOS_ID
	ans_candidate_final[:, Len_candidates[idx] + 1:] = PAD_ID

	for start_block_idx in range(0, Len_candidates[idx].item(), block_len):
	end_block_idx = min(start_block_idx + block_len, Len_candidates[idx].item())

	block_region = torch.zeros(size=(64,), device=device).bool()
	block_region[start_block_idx:end_block_idx] = True

	while True:
	masks_still_to_open = (ans_candidate_final[0] == MASK_ID) & block_region
	if not masks_still_to_open.any().item(): break

	ans_candidate_logits = model.decode(ans_candidate_final, KV_cache, KV_mask).float()

	#penalty wrong end
	ans_candidate_logits[..., (MASK_ID, PAD_ID, EOS_ID)] = -float('inf')

	#penalty repeat
	used = ans_candidate_final[0][ans_candidate_final[0] != MASK_ID].unique()
	lu = ans_candidate_logits[0, :, used]
	ans_candidate_logits[0, :, used] = torch.where(lu > 0, lu / repeat_penalty, lu * repeat_penalty)

	#penalty ngram
	ngrams_ans_end, ngrams_ans_start = ngram(ans_candidate_final[0].tolist(), ngram_ans_len)
	penalize_ngram(ngrams_ans_end, ngrams_ans_start, ngram_ans_len, ans_candidate_final, Len_candidates[idx].item(),
	ans_candidate_logits, ngram_penalty, ANSWER_LEN)
	penalize_ngram(ngram_end_user, ngram_start_user, ngram_user_len, ans_candidate_final, Len_candidates[idx].item(),
	ans_candidate_logits, ngram_penalty, ANSWER_LEN)

	# temperature
	ans_candidate_logits[0] /= temperature

	#top k
	min_logit_per_row = ans_candidate_logits.topk(top_k, dim=2).values[:, :, -1:] # [64, 1]
	ans_candidate_logits = ans_candidate_logits.masked_fill(ans_candidate_logits < min_logit_per_row, -float("inf"))

	# probs and word choosing
	ans_candidate_probs = torch.softmax(ans_candidate_logits, 2).squeeze(0) # [1, 64, 12288] -> [64, 12288]

	if greedy:
	tokens_choice = torch.argmax(ans_candidate_probs, dim=1, keepdim=True)
	else:
	tokens_choice = torch.multinomial(ans_candidate_probs, 1) # [64, 1]
	confidence = ans_candidate_probs.gather(-1, tokens_choice).squeeze(1) # [64, 1] -> [64]

	# open only one mask per step
	i = confidence.masked_fill(~masks_still_to_open, -1).argmax().item()
	ans_candidate_final[0, i] = tokens_choice[i, 0]

	# gstar
	gstar(device, model, margin_for_gstar_replace, start_block_idx, end_block_idx, KV_cache, KV_mask, ans_candidate_final, context_ids,
	rounds=block_rounds)
	this_candidate_score = gstar(device, model, margin_for_gstar_replace, 0, Len_candidates[idx].item(), KV_cache, KV_mask, ans_candidate_final,
	context_ids, rounds=global_rounds, is_global=True)
	if this_candidate_score > best_candidate_score:
	best_candidate_score = this_candidate_score
	best_candidate = ans_candidate_final

	return best_candidate.squeeze(0)

	class OrtBackend():

	def __init__(self, enc_path, dec_path):
	self.encoder = onnxruntime.InferenceSession(enc_path, providers=['CPUExecutionProvider'])
	self.decoder = onnxruntime.InferenceSession(dec_path, providers=['CPUExecutionProvider'])

	def encode(self, context_ids):
	# context_ids [1, 256] torch.tensor
	memory, mask = self.encoder.run(None, {'context_ids': context_ids.numpy()})
	return torch.from_numpy(memory), torch.from_numpy(mask)

	def decode(self, answer_ids, memory, mask):
	#
	return torch.from_numpy(self.decoder.run(None, {'memory': memory[:1].numpy(), 'mask': mask[:1].numpy(), 'answer_ids': answer_ids.numpy()})[0])