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 and last 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])