digital-garbageandClaude Sonnet 4.6
55d905e6ee
.clone() on an inference tensor inherits the inference flag even inside inference_mode(False). torch.empty() inside that context creates a true normal tensor; copy_() fills it from the inference tensor without flag inheritance. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
"""Online value function for reward-guided sampling.
|
|
|
|
Trains a small MLP incrementally on user-rated generations.
|
|
At inference, provides ∂reward/∂conditioning as a per-step steering gradient.
|
|
"""
|
|
|
|
import os
|
|
import random
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class OnlineValueFunction(nn.Module):
|
|
MIN_SAMPLES = 5
|
|
BUFFER_SIZE = 100
|
|
BATCH_SIZE = 16
|
|
TRAIN_STEPS = 20
|
|
|
|
def __init__(self, hidden_dim):
|
|
super().__init__()
|
|
self.hidden_dim = hidden_dim
|
|
mid = min(256, hidden_dim // 4)
|
|
self.net = nn.Sequential(
|
|
nn.Linear(hidden_dim, mid),
|
|
nn.SiLU(),
|
|
nn.Linear(mid, 64),
|
|
nn.SiLU(),
|
|
nn.Linear(64, 1),
|
|
)
|
|
self._optimizer = None
|
|
self.buffer_c = [] # compressed conditioning tensors [hidden_dim]
|
|
self.buffer_r = [] # reward floats
|
|
self.n_trained = 0
|
|
|
|
@property
|
|
def optimizer(self):
|
|
if self._optimizer is None:
|
|
self._optimizer = torch.optim.Adam(self.parameters(), lr=2e-3)
|
|
return self._optimizer
|
|
|
|
@staticmethod
|
|
def compress(conditioning):
|
|
"""[1, seq, D] or [seq, D] → [D] mean vector."""
|
|
c = conditioning.float()
|
|
if c.dim() == 3:
|
|
c = c.squeeze(0)
|
|
return c.mean(dim=0)
|
|
|
|
def forward(self, c):
|
|
return self.net(c)
|
|
|
|
def train_on(self, conditioning, reward):
|
|
"""Add a new (conditioning, reward) sample and do a few training steps."""
|
|
with torch.inference_mode(False), torch.enable_grad():
|
|
c = self.compress(conditioning.clone()).detach().cpu()
|
|
self.buffer_c.append(c)
|
|
self.buffer_r.append(float(reward))
|
|
if len(self.buffer_c) > self.BUFFER_SIZE:
|
|
self.buffer_c = self.buffer_c[-self.BUFFER_SIZE:]
|
|
self.buffer_r = self.buffer_r[-self.BUFFER_SIZE:]
|
|
if len(self.buffer_c) < 2:
|
|
return
|
|
device = next(self.parameters()).device
|
|
for _ in range(self.TRAIN_STEPS):
|
|
idx = random.sample(range(len(self.buffer_c)), min(self.BATCH_SIZE, len(self.buffer_c)))
|
|
c_b = torch.stack([self.buffer_c[i] for i in idx]).to(device)
|
|
r_b = torch.tensor([self.buffer_r[i] for i in idx], device=device).unsqueeze(1)
|
|
self.optimizer.zero_grad()
|
|
F.mse_loss(self.forward(c_b), r_b).backward()
|
|
self.optimizer.step()
|
|
self.n_trained += 1
|
|
|
|
def gradient(self, conditioning):
|
|
"""∂reward/∂conditioning — same shape as input."""
|
|
mlp_device = next(self.parameters()).device
|
|
orig_device = conditioning.device
|
|
with torch.inference_mode(False), torch.enable_grad():
|
|
# new_empty inside inference_mode(False) creates a normal (non-inference) tensor.
|
|
# copy_ fills it from the inference tensor without inheriting the flag.
|
|
c_fresh = torch.empty(conditioning.shape, dtype=torch.float32, device=mlp_device)
|
|
c_fresh.copy_(conditioning)
|
|
c_fresh.requires_grad_(True)
|
|
reward = self.forward(self.compress(c_fresh).unsqueeze(0))
|
|
reward.backward()
|
|
return c_fresh.grad.to(orig_device, conditioning.dtype)
|
|
|
|
def is_ready(self):
|
|
return len(self.buffer_c) >= self.MIN_SAMPLES
|
|
|
|
def save(self, path):
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
buf = torch.stack(self.buffer_c) if self.buffer_c else torch.zeros(0, self.hidden_dim)
|
|
torch.save({
|
|
"hidden_dim": self.hidden_dim,
|
|
"state_dict": self.state_dict(),
|
|
"buffer_c": buf,
|
|
"buffer_r": list(self.buffer_r),
|
|
"n_trained": self.n_trained,
|
|
}, path)
|
|
|
|
@classmethod
|
|
def load(cls, path):
|
|
data = torch.load(path, map_location="cpu", weights_only=False)
|
|
vf = cls(hidden_dim=data["hidden_dim"])
|
|
vf.load_state_dict(data["state_dict"])
|
|
buf = data["buffer_c"]
|
|
vf.buffer_c = [buf[i] for i in range(len(buf))]
|
|
vf.buffer_r = list(data["buffer_r"])
|
|
vf.n_trained = data.get("n_trained", 0)
|
|
return vf
|
|
|
|
@classmethod
|
|
def load_or_create(cls, path, hidden_dim=None):
|
|
if os.path.exists(path):
|
|
try:
|
|
return cls.load(path)
|
|
except Exception:
|
|
pass
|
|
return cls(hidden_dim=hidden_dim) if hidden_dim is not None else None
|