Cache bots instead of game

app.py

@@ -5,19 +5,11 @@ from chessmodels import DCMinMax
 import streamlit as st
 from streamlit_image_coordinates import streamlit_image_coordinates
 
-class HvB(HumanVsBot):
-    @staticmethod
-    @st.cache_resource()
-    def get_hvb_manager(
-        _model,
-        _environment,
-        _extra_model_environment_context,
-        **kwargs):
-        return HvB(
-            model = _model,
-            environment = _environment,
-            extra_model_environment_context = _extra_model_environment_context,
-            **kwargs)
+BOT_SELECITONS = ["DeepChessReplicationMinMax"]
+
+@st.cache_resource
+def load_model(bot_select):
+    return DCMinMax.from_pretrained("mzimm003/{}".format(bot_select))
 
 @st.cache_resource
 def track_play():
@@ -46,10 +38,10 @@ def reset(t_p):
     st.cache_resource.clear()
 
 def play(bot_select, t_p):
-    play = HvB.get_hvb_manager(
-        _model=DCMinMax.from_pretrained("mzimm003/{}".format(bot_select)),
-        _environment=Chess(render_mode="rgb_array"),
-        _extra_model_environment_context=lambda env: {"board":env.board}
+    play = HumanVsBot(
+        model=load_model(bot_select=bot_select),
+        environment=Chess(render_mode="rgb_array"),
+        extra_model_environment_context=lambda env: {"board":env.board},
     )
     col1, col2 = st.columns([3,1])
     with col1:
@@ -81,7 +73,7 @@ def main(kwargs=None):
         with placeholder.container():
             t_p.bot_select = st.selectbox(
                 label="bot",
-                options=["DeepChessReplicationMinMax"],
+                options=BOT_SELECITONS,
             )
             t_p.active_game = st.button("Play bot!")
             if t_p.active_game:

launch.py

@@ -1,4 +1,4 @@
 from streamlit.testing.v1 import AppTest
 
-at = AppTest.from_file("/app/app.py", default_timeout=100)
+at = AppTest.from_file("app.py", default_timeout=100)
 at.run()