Update app.py

app.py

@@ -48,6 +48,17 @@ class UxState(str, Enum):
 # Command-line arguments to control some stuff for easier local testing.
 # Eventually may want to move everything into functions and have a
 # if __name__ == "main" setup instead of everything inline.
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--dry-run", action="store_true",
+    help="Skip sending train request to backend server.",
+)
+parser.add_argument(
+    "--train-endpoint-url", default=None,
+    help="URL of backend server to send train request to. If None, use hardcoded banana setup.",
+)
+cli_args = parser.parse_args()
+
 def setup_session_state():
     if "key" not in st.session_state:
         st.session_state["key"] = uuid.uuid4().hex
@@ -190,10 +201,22 @@ def zip_and_upload_images(identifier: str, uploaded_files: List[BytesIO], image_
         return s3_path
 
 def train_model(model_inputs):
-    api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
-    model_key = "bd2c55f5-84bb-40f9-82fb-196ca68b1c1d"
-    st.markdown(str(model_inputs))
-    _ = banana.run(api_key, model_key, model_inputs)
+    if cli_args.dry_run:
+        logger.info("Skipping model training since --dry-run is enabled.")
+        logger.info(f"model_inputs: {model_inputs}")
+        return
+
+    if cli_args.train_endpoint_url is None:
+        # Use banana backend
+        api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
+        model_key = "1a3b4ce5-164f-4efb-9f4a-c2ad3a930d0b"
+        st.markdown(str(model_inputs))
+        print(model_inputs)
+        _ = banana.run(api_key, model_key, model_inputs)
+    else:
+        # Send request directly to specified url
+        _ = requests.post(cli_args.train_endpoint_url, json=model_inputs)
+
 
 def switch_ux_state(new_state: UxState):
     st.session_state['ux_state'] = new_state