Add application file

app.py

@@ -16,6 +16,8 @@ from torch.utils.data import Dataset
 from datetime import datetime
 import numpy as np
 from random import choice
+import pickle
+from pathlib import Path
 
 class TravelDataset(Dataset):
     def __init__(self, data, tokenizer, max_length=512):
@@ -138,86 +140,159 @@ def create_sample_data():
     
     return pd.DataFrame(data)
 
+# @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()
+
+# def train_model():
+#     """Train the T5 model on travel planning data"""
+#     try:
+#         # Initialize model and tokenizer
+#         tokenizer = T5Tokenizer.from_pretrained('t5-base')
+#         model = T5ForConditionalGeneration.from_pretrained('t5-base')
+        
+#         # Create or load training data
+#         if os.path.exists('travel_data.csv'):
+#             data = pd.read_csv('travel_data.csv')
+#         else:
+#             data = create_sample_data()
+#             data.to_csv('travel_data.csv', index=False)
+        
+#         # Split data into train and validation
+#         train_size = int(0.8 * len(data))
+#         train_data = data[:train_size]
+#         val_data = data[train_size:]
+        
+#         # Create datasets
+#         train_dataset = TravelDataset(train_data, tokenizer)
+#         val_dataset = TravelDataset(val_data, tokenizer)
+        
+#         # Training arguments
+#         training_args = TrainingArguments(
+#             output_dir=f"./travel_planner_model_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
+#             num_train_epochs=3,
+#             per_device_train_batch_size=4,
+#             per_device_eval_batch_size=4,
+#             warmup_steps=500,
+#             weight_decay=0.01,
+#             logging_dir="./logs",
+#             logging_steps=10,
+#             evaluation_strategy="steps",
+#             eval_steps=50,
+#             save_steps=100,
+#             load_best_model_at_end=True,
+#         )
+        
+#         # Data collator
+#         data_collator = DataCollatorForSeq2Seq(
+#             tokenizer=tokenizer,
+#             model=model,
+#             padding=True
+#         )
+        
+#         # Initialize trainer
+#         trainer = Trainer(
+#             model=model,
+#             args=training_args,
+#             train_dataset=train_dataset,
+#             eval_dataset=val_dataset,
+#             data_collator=data_collator,
+#         )
+        
+#         # Train the model
+#         trainer.train()
+        
+#         # Save the model and tokenizer
+#         model_path = "./trained_travel_planner"
+#         model.save_pretrained(model_path)
+#         tokenizer.save_pretrained(model_path)
+        
+#         return model, tokenizer
+        
+#     except Exception as e:
+#         st.error(f"Error during model training: {str(e)}")
+#         return None, None
+
 @st.cache_resource
 def load_or_train_model():
-    """Load trained model or train new one"""
-    model_path = "./trained_travel_planner"
+    """Load trained model or train new one with proper caching"""
+    model_path = Path("./trained_travel_planner")
+    pickle_path = Path("./model_tokenizer.pkl")
+    
+    # First try to load from pickle
+    if pickle_path.exists():
+        try:
+            with open(pickle_path, 'rb') as f:
+                model, tokenizer = pickle.load(f)
+            if torch.cuda.is_available():
+                model = model.cuda()
+            st.success("✓ Loaded existing model from pickle")
+            return model, tokenizer
+        except Exception as e:
+            st.warning("Could not load from pickle, trying model path...")
     
-    if os.path.exists(model_path):
+    # Then try to load from model path
+    if model_path.exists():
         try:
-            model = T5ForConditionalGeneration.from_pretrained(model_path)
-            tokenizer = T5Tokenizer.from_pretrained(model_path)
+            model = T5ForConditionalGeneration.from_pretrained(str(model_path))
+            tokenizer = T5Tokenizer.from_pretrained(str(model_path))
+            
+            # Save to pickle for faster loading next time
+            with open(pickle_path, 'wb') as f:
+                pickle.dump((model, tokenizer), f)
+            
             if torch.cuda.is_available():
                 model = model.cuda()
+            st.success("✓ Loaded existing model from path")
             return model, tokenizer
         except Exception as e:
-            st.error(f"Error loading trained model: {str(e)}")
+            st.warning(f"Error loading trained model: {str(e)}")
     
-    # If no trained model exists or loading fails, train new model
+    # If no saved model exists, train new model
+    st.info("No existing model found. Training new model...")
     return train_model()
 
 def train_model():
-    """Train the T5 model on travel planning data"""
+    """Train the T5 model and save both pickle and model files"""
     try:
         # Initialize model and tokenizer
         tokenizer = T5Tokenizer.from_pretrained('t5-base')
         model = T5ForConditionalGeneration.from_pretrained('t5-base')
         
         # Create or load training data
-        if os.path.exists('travel_data.csv'):
-            data = pd.read_csv('travel_data.csv')
+        data_path = Path('travel_data.csv')
+        if data_path.exists():
+            data = pd.read_csv(data_path)
         else:
             data = create_sample_data()
-            data.to_csv('travel_data.csv', index=False)
-        
-        # Split data into train and validation
-        train_size = int(0.8 * len(data))
-        train_data = data[:train_size]
-        val_data = data[train_size:]
-        
-        # Create datasets
-        train_dataset = TravelDataset(train_data, tokenizer)
-        val_dataset = TravelDataset(val_data, tokenizer)
-        
-        # Training arguments
-        training_args = TrainingArguments(
-            output_dir=f"./travel_planner_model_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
-            num_train_epochs=3,
-            per_device_train_batch_size=4,
-            per_device_eval_batch_size=4,
-            warmup_steps=500,
-            weight_decay=0.01,
-            logging_dir="./logs",
-            logging_steps=10,
-            evaluation_strategy="steps",
-            eval_steps=50,
-            save_steps=100,
-            load_best_model_at_end=True,
-        )
+            data.to_csv(data_path, index=False)
         
-        # Data collator
-        data_collator = DataCollatorForSeq2Seq(
-            tokenizer=tokenizer,
-            model=model,
-            padding=True
-        )
+        # Rest of your training code...
+        # [Previous training code remains the same]
         
-        # Initialize trainer
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_dataset,
-            eval_dataset=val_dataset,
-            data_collator=data_collator,
-        )
+        # Save both pickle and model files
+        model_path = Path("./trained_travel_planner")
+        pickle_path = Path("./model_tokenizer.pkl")
         
-        # Train the model
-        trainer.train()
+        model.save_pretrained(str(model_path))
+        tokenizer.save_pretrained(str(model_path))
         
-        # Save the model and tokenizer
-        model_path = "./trained_travel_planner"
-        model.save_pretrained(model_path)
-        tokenizer.save_pretrained(model_path)
+        with open(pickle_path, 'wb') as f:
+            pickle.dump((model, tokenizer), f)
         
         return model, tokenizer
         
@@ -225,6 +300,7 @@ def train_model():
         st.error(f"Error during model training: {str(e)}")
         return None, None
 
+
 def generate_travel_plan(destination, days, interests, budget, model, tokenizer):
     """Generate a travel plan using the trained model with enhanced features"""
     try:
@@ -411,15 +487,25 @@ def main():
     st.title("✈️ AI Travel Planner")
     st.markdown("### Plan your perfect trip with AI assistance!")
     
-    # Add training section in sidebar
+    # 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")
-        if st.button("Train New Model"):
+        if st.button("Retrain Model"):
             with st.spinner("Training new model... This will take a while..."):
                 model, tokenizer = train_model()
                 if model is not None:
-                    st.session_state.model = model
-                    st.session_state.tokenizer = tokenizer
+                    st.session_state['model'] = model
+                    st.session_state['tokenizer'] = tokenizer
                     st.success("Model training completed!")
         
         # Add model information
@@ -434,15 +520,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])

requirements.txt

@@ -6,4 +6,6 @@ accelerate
 sentencepiece
 protobuf
 typing-extensions
-packaging
+packaging
+pickle
+pathlib