Views
No views yet
std_math prediction in data mixture optimization.K(x, x') = exp(-γ · ||x - x'||₁)[if_prop1, math_prop1, math_prop2] (values in [0, 1])1import torch
2from huggingface_hub import hf_hub_download
3from safetensors.torch import load_file
4
5
6class KernelRegressionModel(torch.nn.Module):
7 def __init__(self, dual_coef, X_fit, gamma=0.1):
8 super().__init__()
9 self.gamma = gamma
10 self.register_buffer("dual_coef", dual_coef)
11 self.register_buffer("X_fit", X_fit)
12
13 def forward(self, x):
14 dist = torch.cdist(x, self.X_fit, p=1)
15 K = torch.exp(-self.gamma * dist)
16 return K @ self.dual_coef
17
18
19path = hf_hub_download("chewwt/dm_qwen4b_noise_emulator", "noise_model.safetensors")
20tensors = load_file(path)
21model = KernelRegressionModel(tensors["dual_coef"], tensors["X_fit"])
22model.eval()
23
24# x: (batch, 3) float64 tensor, features in [0, 1]
25x = torch.tensor([[0.3, 0.4, 0.2]], dtype=torch.float64)
26with torch.no_grad():
27 sigma = model(x) # predicted std_math