使用 Spring AI 调用本地 模型实现

在本篇博客中，我们将学习如何使用 Spring AI 框架调用本地的 PyTorch 模型，并通过 Spring Boot 提供一个预测接口。Spring AI 是一个用于将人工智能应用集成到 Spring 生态系统中的框架，它支持多种 AI 模型和数据源的集成，帮助开发者将 AI 模型无缝地集成到 Java 应用中。

1. 准备 PyTorch 模型

首先，我们需要训练并保存一个 PyTorch 模型。这里我们使用一个简单的神经网络模型作为示例。训练并保存模型后，我们会将其转换为 TorchScript 格式，TorchScript 是 PyTorch 提供的一种中间表示格式，可以在 C++ 和 Java 环境中使用。

以下是一个简单的 PyTorch 模型示例：

import torch
import torch.nn as nn

# 示例模型（简单的神经网络）
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

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

# 加载模型
model = SimpleModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()

# 保存模型为 TorchScript 格式（可用于 Java）
traced_model = torch.jit.trace(model, torch.randn(1, 10))
traced_model.save("model_traced.pt")

运行这段代码，你将得到一个 model_traced.pt 文件，该文件将用于后续的 Spring AI 集成。

2. 集成 PyTorch 模型到 Spring Boot 项目