Add application file

app.py

@@ -17,6 +17,17 @@ from datetime import datetime
 import numpy as np
 from random import choice
 
+# Add these settings at the top level of your script
+st.set_page_config(
+    page_title="AI Travel Planner",
+    page_icon="✈️",
+    layout="wide"
+)
+
+# Initialize session state for model
+if 'initialized' not in st.session_state:
+    st.session_state.initialized = False
+
 class TravelDataset(Dataset):
     def __init__(self, data, tokenizer, max_length=512):
         """
@@ -138,23 +149,54 @@ def create_sample_data():
     
     return pd.DataFrame(data)
 
-@st.cache_resource
+@st.cache_resource(show_spinner=False)
 def load_or_train_model():
     """Load trained model or train new one"""
+    # Check if model exists in session state first
+    if hasattr(st.session_state, 'model') and hasattr(st.session_state, 'tokenizer'):
+        return st.session_state.model, st.session_state.tokenizer
+        
     model_path = "./trained_travel_planner"
     
     if os.path.exists(model_path):
         try:
+            # Load existing model and tokenizer
             model = T5ForConditionalGeneration.from_pretrained(model_path)
             tokenizer = T5Tokenizer.from_pretrained(model_path)
             if torch.cuda.is_available():
                 model = model.cuda()
+            # Store in session state
+            st.session_state.model = model
+            st.session_state.tokenizer = tokenizer
             return model, tokenizer
         except Exception as e:
-            st.error(f"Error loading trained model: {str(e)}")
+            st.error(f"Error loading model: {str(e)}")
+    
+    # If no model exists, train new one
+    st.warning("No trained model found. Training new model...")
+    model, tokenizer = train_model()
+    # Store in session state
+    st.session_state.model = model
+    st.session_state.tokenizer = tokenizer
+    return model, tokenizer
+
+# @st.cache_resource
+# def load_or_train_model():
+#     """Load trained model or train new one"""
+#     model_path = "./trained_travel_planner"
+    
+#     if os.path.exists(model_path):
+#         try:
+#             model = T5ForConditionalGeneration.from_pretrained(model_path)
+#             tokenizer = T5Tokenizer.from_pretrained(model_path)
+#             if torch.cuda.is_available():
+#                 model = model.cuda()
+#             return model, tokenizer
+#         except Exception as e:
+#             st.error(f"Error loading trained model: {str(e)}")
     
-    # If no trained model exists or loading fails, train new model
-    return train_model()
+#     # If no trained model exists or loading fails, train new model
+#     return train_model()
 
 def train_model():
     """Train the T5 model on travel planning data"""
@@ -411,15 +453,21 @@ def main():
     st.title("✈️ AI Travel Planner")
     st.markdown("### Plan your perfect trip with AI assistance!")
     
-    # Load model only if not in session state
-    if 'model' not in st.session_state or 'tokenizer' not in st.session_state:
+    # Move model loading to a initialization section
+    if 'initialized' not in st.session_state:
         with st.spinner("Loading AI model... Please wait..."):
             model, tokenizer = load_or_train_model()
-            if model is None or tokenizer is None:
-                st.error("Failed to load/train the AI model. Please try again.")
-                return
-            st.session_state['model'] = model
-            st.session_state['tokenizer'] = tokenizer
+            st.session_state.initialized = True
+    
+    # # Load model only if not in session state
+    # if 'model' not in st.session_state or 'tokenizer' not in st.session_state:
+    #     with st.spinner("Loading AI model... Please wait..."):
+    #         model, tokenizer = load_or_train_model()
+    #         if model is None or tokenizer is None:
+    #             st.error("Failed to load/train the AI model. Please try again.")
+    #             return
+    #         st.session_state['model'] = model
+    #         st.session_state['tokenizer'] = tokenizer
     
     # Add training button in sidebar only
     with st.sidebar: