| |
| |
| |
| class AdaptiveGate(nn.Module): |
| """ |
| y_final = w_f * Retrieve(S_best) + (1 - w_f) * Generate(z_0) |
| Var(y_final) = (1 - w_f)^2 * sigma_z0^2 |
| As w_f -> 1, Var -> 0: hallucination prevention. |
| """ |
| def __init__(self, d_model: int, *, rngs: nn.Rngs): |
| super().__init__() |
| self.d_model = d_model |
| self.sigma_raw = nn.Param(jnp.array(0.0)) |
|
|
| @property |
| def sigma(self): |
| return jax.nn.softplus(self.sigma_raw.value) |
|
|
| def __call__(self, retrieved, generated, w_f): |
| if w_f.ndim < retrieved.ndim: |
| w_f = w_f[..., None] |
| y_final = w_f * retrieved + (1.0 - w_f) * generated |
| var_y_final = (1.0 - w_f)**2 * self.sigma**2 |
| return y_final, var_y_final |
|
|
| print("[5/6] AdaptiveGate defined") |
|
|
