import os

os.environ.setdefault("KERAS_BACKEND", "jax")

from pathlib import Path

import keras
import numpy as np


def build_model():
    return keras.Sequential(
        [
            keras.Input(
                shape=(2,),
                name="features",
            ),
            keras.layers.Dense(
                1,
                name="score",
            ),
        ]
    )