Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,8 @@
 import streamlit as st
 import torch
+import os
 from transformers import AutoTokenizer
-from predict_utils import predict_injury  # assumes predict_injury is defined here
+from predict_utils import predict_injury  # assumes predict_injury is defined in this module
 
 # ✅ Load tokenizer
 @st.cache_resource
@@ -11,32 +12,36 @@ def load_tokenizer():
 # ✅ Load model
 @st.cache_resource
 def load_model():
-    model_path = "model/final_injury_model.pt"  # or use a Hugging Face model ID if from hub
+    model_path = "model/final_injury_model.pt"  # adjust if needed
+    if not os.path.exists(model_path):
+        st.error(f"Model file not found at: {model_path}")
+        return None
+
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = torch.load(model_path, map_location=device)
-    if not os.path.exists(model_path):
-    st.error(f"Model file not found at: {model_path}")
-    return None
     model.eval()
- 
-
     return model
+
+# ✅ Main Streamlit app
 def main():
-    st.title("NBA Injury Type & Duration Classifier 🏀")
+    st.title("🏀 NBA Injury Type & Duration Classifier")
 
     model = load_model()
     tokenizer = load_tokenizer()
 
-    st.markdown("Enter injury note and details to get predictions")
+    if model is None:
+        st.stop()  # prevent rest of UI from running if model isn't loaded
+
+    st.markdown("Enter an injury report and structured data to get predictions.")
 
-    # form inputs
-    text = st.text_input("Injury description", "player has a sprained ankle")
-    prior_injuries = st.number_input("Prior Injuries", min_value=0, value=1)
-    injury_type_id = st.selectbox("Injury Type", {"bone": 0, "muscle": 1, "joint": 2, "ligament": 3})
-    position_id = st.selectbox("Position", {"PG": 1, "SG": 2, "SF": 3, "PF": 4, "C": 5})
+    # user inputs
+    text = st.text_input("📝 Injury Description", "player has a sprained ankle")
+    prior_injuries = st.number_input("🔁 Prior Injuries", min_value=0, value=1)
+    injury_type_id = st.selectbox("🩻 Injury Type", {"bone": 0, "muscle": 1, "joint": 2, "ligament": 3})
+    position_id = st.selectbox("🏀 Player Position", {"PG": 1, "SG": 2, "SF": 3, "PF": 4, "C": 5})
 
     # prediction
-    if st.button("Predict"):
+    if st.button("🔮 Predict"):
         label_map_type = ['bone', 'muscle', 'joint', 'ligament']
         label_map_duration = ['short', 'medium', 'long']
 
@@ -51,7 +56,8 @@ def main():
             label_map_duration=label_map_duration
         )
 
-        st.success(f"**Injury Type:** {type_label} ({type_conf:.1%})")
-        st.success(f"**Recovery Duration:** {duration_label} ({duration_conf:.1%})")
+        st.success(f"**Predicted Injury Type:** `{type_label}` ({type_conf:.1%} confidence)")
+        st.success(f"**Predicted Duration:** `{duration_label}` ({duration_conf:.1%} confidence)")
+
 if __name__ == "__main__":
     main()