# 项目文件：model_project/src/modeldelivery/model.py
from torch import nn


class ReadingRegressor(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)  # 沿用上一章两个特征到一个目标的结构。

    def forward(self, x):
        return self.linear(x)

# 项目文件：model_project/src/modeldelivery/training.py
import math

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from .model import ReadingRegressor


def fit(data, epochs=35, seed=7):
    if epochs < 1:
        raise ValueError("训练轮数需要大于零")
    # 固定 CPU 计算和随机来源，便于在新进程中核对同一轮训练。
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    torch.use_deterministic_algorithms(True)
    x, y = torch.from_numpy(data.x), torch.from_numpy(data.y)
    train_ids, val_ids = data.splits["train"], data.splits["val"]
    loader = DataLoader(
        TensorDataset(x[train_ids], y[train_ids]),
        batch_size=24,
        shuffle=True,
        generator=torch.Generator().manual_seed(seed + 1),
    )
    model, loss_fn = ReadingRegressor(), nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.03)
    best_loss, best_epoch, best, history = float("inf"), 0, None, []

# 续写同一文件：model_project/src/modeldelivery/training.py
    # 训练批次更新参数，验证只读取并保存较优快照。
    for epoch in range(epochs):
        model.train()
        total = 0.0
        for features, target in loader:
            optimizer.zero_grad(set_to_none=True)
            loss = loss_fn(model(features), target)
            if not torch.isfinite(loss):
                raise ValueError("训练损失出现非有限数值，请检查这批输入")
            loss.backward()
            optimizer.step()
            total += loss.item() * len(features)
        model.eval()
        with torch.inference_mode():
            validation = loss_fn(model(x[val_ids]), y[val_ids]).item()
        history.append(
            {
                "epoch": epoch + 1,
                "train_mse": total / len(train_ids),
                "val_mse": validation,
            }
        )
        if validation < best_loss:
            best_loss, best_epoch = validation, epoch + 1
            best = {
                name: value.detach().cpu().clone()
                for name, value in model.state_dict().items()
            }

# 续写同一文件：model_project/src/modeldelivery/training.py
    # 轮数选择结束后，才用选定快照读取测试组。
    model.load_state_dict(best)
    test_ids = data.splits["test"]
    with torch.inference_mode():
        test_mse = loss_fn(model(x[test_ids]), y[test_ids]).item()
    scale = data.scaler["y_scale"].item()
    baseline = math.sqrt(float((data.y[test_ids] ** 2).mean())) * scale
    metrics = {
        "best_epoch": best_epoch,
        "val_rmse": math.sqrt(best_loss) * scale,
        "test_rmse": math.sqrt(test_mse) * scale,
        "baseline_rmse": baseline,
        "counts": {name: len(ids) for name, ids in data.splits.items()},
    }
    return best, metrics, history