import numpy as np

import torch

from tqdm import tqdm as tq

from mucola import MuCoLa

BOS_ID = 764

DISC_EPS = 0.3

EPS_P = 1.5

class SVS(MuCoLa):

def __init__(self, c_model_info, args, pad_token_id, target_dict=None, is_hmc=False):

super().__init__(c_model_info, args, pad_token_id, target_dict)

self.is_baseline = is_hmc

def energy_e_function(self, e2, max_len, context_len=None, target=0):

e_idx, e = e2

h = \

self.g_model.forward(inputs_embeds=e[:, :-1, :], output_hidden_states=True)[

'hidden_states'][

-1][:, context_len - 1:, :]

word_embeddings = self.g_model.transformer.wte.weight

denom = h @ word_embeddings.T

denom = torch.logsumexp(denom, dim=-1)

remained_e = (

e[:, context_len:] - self.g_model.transformer.wpe.weight[

context_len:max_len + context_len, :].unsqueeze(0))

nom = torch.sum(remained_e * h, -1)

loss = - nom + denom

c_energy = self.c_model.energy(e, target)

return loss + c_energy

def find_intersec_point(self, e_t, e_s, p, eps_p):

n = e_t - e_s

mid = (e_t + e_s) / 2.

intersec_p = (torch.dot(n, mid) - torch.dot(n, e_s)) / torch.dot(n, p)

if not intersec_p or (abs(intersec_p) / eps_p) > 1.:

print(abs(intersec_p) / eps_p)

print("ERROR in discontinuity")

return False, None

return True, intersec_p

def decompose(self, p, ref):

p_perp = (torch.dot(p, ref) / torch.dot(ref, ref)) * ref

p_par = p - p_perp

return p_perp, p_par

def refract_reflect(self, p, ref, delta_en, context_len=None):

p_perp, p_par = self.decompose(p, ref)

perp_norm = torch.dot(p_perp, p_perp)

if perp_norm > 2 * delta_en:

message = "refract"

# print("refract")

p_perp = torch.sqrt(perp_norm - 2 * delta_en) * (p_perp / torch.sqrt(perp_norm))

else:

message = "reflect"

# print("reflect")

p_perp = -1 * p_perp

p = p_par + p_perp

return p

def find_discontinuity(self, e_s, e_t, p, eps_p, context=None, context_len=None,

target=0):

with torch.no_grad():

s_idx, e_s_proj = self.project_embeds(e_s, context, context_len)

t_idx, e_t_proj = self.project_embeds(e_t, context, context_len)

delta_en = self.energy_e_function((t_idx, e_t_proj), self.max_len,

context_len,

target,

).mean() - \

self.energy_e_function((s_idx, e_s_proj), self.max_len,

context_len,

target,

).mean()

int_p = torch.zeros(p.shape[0]).to(self.device) + DISC_EPS

if delta_en > 1e-10:

for j in range(self.max_len):

i = j + context_len

is_disc, found_p = self.find_intersec_point(e_t[i], e_s[i], p[i], eps_p)

if is_disc:

int_p[i] = found_p

p[i] = self.refract_reflect(p[i].clone(), e_t[i] - e_s[i], delta_en,

context_len)

return int_p

def simulate_rhmc(self, e, p, eps_p, context=None, context_len=None, target=0):

t = 0

e = e.squeeze(0)

p = p.squeeze(0)

for _ in range(int(1. / DISC_EPS)):

old_p = p.clone()

e_new = e.clone() + eps_p * DISC_EPS * old_p

delta_t = self.find_discontinuity(e, e_new, p, eps_p, context, context_len,

target)

delta_t = delta_t.unsqueeze(-1)

t = torch.clamp(t + delta_t, max=1.)

e_new = e.clone() + eps_p * delta_t * old_p

e = e_new

return e.unsqueeze(0)

def leapfrog_update(self, step_size, e_proj, e, p, eps_p, baseline=True, context=None,

context_len=None, target=0):

for _ in range(1):

e_proj = torch.nn.Parameter(e_proj, requires_grad=True)

energy = self.energy_e_function((e_proj, e_proj), self.max_len,

context_len,

target).mean()

energy.backward(retain_graph=True)

p_mid = p - (step_size / 2) * (

e + e_proj.grad / torch.sqrt(e_proj.grad * e_proj.grad))

if baseline:

e_next = e + eps_p * p_mid

else:

e_next = self.simulate_rhmc(e, p_mid, eps_p, context, context_len, target)

if e_next is None:

return None

with torch.no_grad():

e_next_proj_idx, e_next_proj = self.project_embeds(e_next.data.detach(),

context,

context_len)

e_proj, e = e_next_proj.data.detach(), e_next

return e_proj, e_next_proj_idx, e

def log_results(self, projected_idx):

e = projected_idx[0]

sample = self.tokenizer.decode(e, skip_special_tokens=True)

with torch.no_grad():

ppl = self.g_model(input_ids=e.unsqueeze(0), labels=e.unsqueeze(0))[0]

ppl = np.exp(ppl.detach().cpu().numpy())

print("[{:.2f}]: {}".format(ppl, sample))

return sample, ppl

def predict_with_control(self, save_dir, name, targets=None, contexts=None):

if not self.controlled:

targets = [0] # dummy variable

else:

assert targets is not None # for controlled sampling targets must be provided

if contexts is None:

contexts = [torch.tensor([BOS_ID]).unsqueeze(0).to(self.device)]

for context in contexts:

for target in targets:

if self.controlled:

print(f"{'=' * 10}target: {self.target_dict[target]} {'=' * 10}")

for _ in tq(range(self.number_of_samples)):

eps_p = EPS_P

tol = 25

last_energy = 0.0

e, _ = self.initialize(context=context)

context_len = context.shape[1]

e_proj = e.clone()

best_e = e_proj

best_en = 1e10

pbar = tq(range(self.steps))

with torch.enable_grad():

for i in pbar:

p = torch.normal(mean=0.0, std=1., size=e.size()).to(

self.device) # 1 x max_len x d

res = self.leapfrog_update(self.step_size,

e_proj.clone(), e.clone(), p.clone(),

eps_p,

baseline=self.is_baseline,

context=context,

context_len=context_len,

target=target)

if res is None:

continue

e_next_proj, e_next_idx, e_next = res

p_accept = 1.

eps_p = self.update_noise_variance(eps_p)

coin = np.random.rand()

if coin < p_accept:

e_proj, e, e_idx = e_next_proj, e_next, e_next_idx

if i % self.inner_steps == 0:

self.log_results(e_idx)

energy = self.energy_e_function((e, e_proj), self.max_len,

context_len,

target).mean()

if (last_energy - energy.mean()) < 1e-5:

tol -= 1

else:

tol = 25

if tol == 0:

break

last_energy = energy.mean()

if energy.mean() < best_en:

best_en = energy.mean()

best_e = e_idx.clone()

pbar.set_description(f"e {energy.mean()}", refresh=True)

pbar.update()

sample, ppl = self.log_results(best_e)

self.update_results_dict(sample, ppl, 0, 0, 0)

self.save_results_dict(save_dir, name)

self.save_results_dict(save_dir, name)