from pathlib import Path

import torch
from torch import nn


class ScoreModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 2),
        )

    def forward(self, features):
        return torch.softmax(self.layers(features), dim=-1)


torch.manual_seed(7)

model = ScoreModel().eval()
example_input = torch.randn(2, 4)

with torch.inference_mode():
    eager_output = model(example_input)

scripted = torch.jit.trace(model, example_input)
artifact = Path("score_model_torchscript.pt")
scripted.save(artifact)

loaded = torch.jit.load(artifact, map_location="cpu")
test_input = torch.randn(3, 4)

with torch.inference_mode():
    loaded_output = loaded(test_input)
    eager_check = model(test_input)

print(f"torch version: {torch.__version__}")
print("export method: torch.jit.trace")
print(f"scripted type: {type(scripted).__name__}")
print(f"eager output shape: {tuple(eager_output.shape)}")
print(f"artifact: {artifact} ({artifact.stat().st_size} bytes)")
print(f"loaded output shape: {tuple(loaded_output.shape)}")
print(f"outputs match eager: {torch.allclose(loaded_output, eager_check)}")