Add application file

app.py

@@ -16,8 +16,6 @@ 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):
@@ -140,159 +138,86 @@ 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 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...")
+    """Load trained model or train new one"""
+    model_path = "./trained_travel_planner"
     
-    # Then try to load from model path
-    if model_path.exists():
+    if os.path.exists(model_path):
         try:
-            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)
-            
+            model = T5ForConditionalGeneration.from_pretrained(model_path)
+            tokenizer = T5Tokenizer.from_pretrained(model_path)
             if torch.cuda.is_available():
                 model = model.cuda()
-            st.success("✓ Loaded existing model from path")
             return model, tokenizer
         except Exception as e:
-            st.warning(f"Error loading trained model: {str(e)}")
+            st.error(f"Error loading trained model: {str(e)}")
     
-    # If no saved model exists, train new model
-    st.info("No existing model found. Training new model...")
+    # If no trained model exists or loading fails, train new model
     return train_model()
 
 def train_model():
-    """Train the T5 model and save both pickle and model files"""
+    """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
-        data_path = Path('travel_data.csv')
-        if data_path.exists():
-            data = pd.read_csv(data_path)
+        if os.path.exists('travel_data.csv'):
+            data = pd.read_csv('travel_data.csv')
         else:
             data = create_sample_data()
-            data.to_csv(data_path, index=False)
+            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,
+        )
         
-        # Rest of your training code...
-        # [Previous training code remains the same]
+        # Data collator
+        data_collator = DataCollatorForSeq2Seq(
+            tokenizer=tokenizer,
+            model=model,
+            padding=True
+        )
         
-        # Save both pickle and model files
-        model_path = Path("./trained_travel_planner")
-        pickle_path = Path("./model_tokenizer.pkl")
+        # Initialize trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=val_dataset,
+            data_collator=data_collator,
+        )
         
-        model.save_pretrained(str(model_path))
-        tokenizer.save_pretrained(str(model_path))
+        # Train the model
+        trainer.train()
         
-        with open(pickle_path, 'wb') as f:
-            pickle.dump((model, tokenizer), f)
+        # Save the model and tokenizer
+        model_path = "./trained_travel_planner"
+        model.save_pretrained(model_path)
+        tokenizer.save_pretrained(model_path)
         
         return model, tokenizer
         
@@ -300,7 +225,6 @@ 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:

requirements.txt

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