Spaces:
Sleeping
Sleeping
File size: 2,756 Bytes
7d455d3 bd51db5 7d455d3 bd51db5 7d455d3 bd51db5 0b6ae5f 7d455d3 bd51db5 7d455d3 bd51db5 7d455d3 bd51db5 7d455d3 bd51db5 7d455d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
import torch
import wandb
from plot2eq.inference.pipeline import predict_top_k_equations
@torch.no_grad()
def create_val_predictions_table(model, points, true_tokens, tokenizer, num_examples=5):
num_examples = min(num_examples, points.size(0))
sample_points = points[:num_examples]
sample_true_tokens = true_tokens[:num_examples]
val_table = wandb.Table(
columns=["Plot", "True Equation", "Predicted Equation", "MSE", "Score"]
)
for i in range(num_examples):
y_vals = sample_points[i, 0].cpu().numpy()
mask = sample_points[i, 1].cpu().numpy().astype(bool)
x_vals = np.linspace(0, 1, len(y_vals))
x_data_fit = x_vals[mask]
y_data_fit = y_vals[mask]
results = predict_top_k_equations(
model,
points_tensor=sample_points[i : i + 1],
tokenizer=tokenizer,
x_data=x_data_fit,
y_data=y_data_fit,
beam_size=5,
top_k=1,
length_penalty=0.01,
)
if results:
res = results[0]
pred_expr = res["best_expr"]
mse = res["mse"]
score = res["score"]
else:
pred_expr = None
mse = float("inf")
score = 0.0
fig, ax = plt.subplots(figsize=(5, 4))
ax.plot(x_data_fit, y_data_fit, color="black", linewidth=3, label="Target")
pred_str = "Parse Error / Fit Failed"
if pred_expr is not None:
pred_str = f"$${sp.latex(pred_expr)}$$"
x_sym = sp.Symbol("x", real=True)
f_expr = sp.lambdify(x_sym, pred_expr, modules=["numpy"])
try:
preds_y = f_expr(x_data_fit)
if np.isscalar(preds_y):
preds_y = np.full_like(x_data_fit, preds_y)
ax.plot(
x_data_fit,
preds_y,
color="red",
linestyle="--",
linewidth=2.5,
label="Predicted",
)
except Exception:
pass
ax.set_ylim(-0.05, 1.05)
ax.legend()
plt.tight_layout()
true_toks = [
t.item() for t in sample_true_tokens[i] if t.item() not in (0, 1, 2)
]
try:
true_expr = tokenizer.token_seq_to_expr(torch.tensor(true_toks))
true_str = f"$${sp.latex(true_expr)}$$"
except Exception:
true_str = "Parse Error"
val_table.add_data(
wandb.Image(fig), true_str, pred_str, f"{mse:.5f}", f"{score:.5f}"
)
plt.close(fig)
return val_table
|