import numpy as np def gs1_adv(vec, prompts, rewards, T=0.1): X = vec.transform(prompts).toarray().astype(np.float64) norms = np.linalg.norm(X, axis=1, keepdims=True) Xn = X / np.where(norms == 0, 1.0, norms) S = Xn @ Xn.T W = np.exp(S / T) np.fill_diagonal(W, 0.0) r = np.asarray(rewards, dtype=np.float64) b = (W @ r) / W.sum(1) return r - b, b