Update app.py

app.py

@@ -17,10 +17,6 @@ from datetime import datetime
 import numpy as np
 from random import choice
 
-# 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):
         """
@@ -142,54 +138,23 @@ def create_sample_data():
     
     return pd.DataFrame(data)
 
-@st.cache_resource(show_spinner=False)
+@st.cache_resource
 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"
+     """Load trained model or train new one"""
+     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 model: {str(e)}")
+     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 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"""
@@ -446,22 +411,6 @@ def main():
     st.title("✈️ AI Travel Planner")
     st.markdown("### Plan your perfect trip with AI assistance!")
     
-    # 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()
-            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:
         st.header("Model Management")
@@ -485,15 +434,15 @@ def main():
             - 5 interest combinations
             """)
     
-    # # Load or train model
-    # if 'model' 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
+     # Load or train model
+     if 'model' 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
     
     # Create two columns for input form
     col1, col2 = st.columns([2, 1])