Update data/prepare_datasets.py

data/prepare_datasets.py

@@ -8,6 +8,9 @@ from sklearn.model_selection import train_test_split
 import hashlib
 import json
 from datetime import datetime
+from data.data_validator import DataValidationPipeline
+from data.validation_schemas import ValidationLevel, DataSource
+from typing import Tuple, Dict
 
 # Configure logging
 logging.basicConfig(
@@ -131,7 +134,7 @@ class DatasetPreparer:
             ], ignore_index=True)
             
             logger.info(f"Combined Kaggle dataset: {len(df_combined)} samples")
-            return df_combined
+            return self.validate_dataset_with_schemas(df_combined, 'kaggle_combined')
             
         except Exception as e:
             logger.error(f"Error loading Kaggle dataset: {e}")
@@ -201,7 +204,7 @@ class DatasetPreparer:
             if liar_dfs:
                 combined_liar = pd.concat(liar_dfs, ignore_index=True)
                 logger.info(f"Combined LIAR dataset: {len(combined_liar)} samples")
-                return combined_liar
+                return self.validate_dataset_with_schemas(combined_liar, 'liar_combined')
             else:
                 logger.warning("No LIAR data could be processed")
                 return None
@@ -226,7 +229,8 @@ class DatasetPreparer:
         # Validate text quality
         valid_mask = df['text'].apply(self.validate_text_quality)
         df = df[valid_mask]
-        logger.info(f"Removed {initial_count - len(valid_mask.sum())} low-quality texts")
+        # logger.info(f"Removed {initial_count - len(valid_mask.sum())} low-quality texts")
+        logger.info(f"Removed {initial_count - valid_mask.sum()} low-quality texts")
         
         # Remove duplicates
         before_dedup = len(df)
@@ -300,63 +304,124 @@ class DatasetPreparer:
         return float(np.mean(scores))
     
     def prepare_datasets(self) -> Tuple[bool, str]:
-        """Main function to prepare and combine all datasets"""
+        """Main method to prepare all datasets with validation"""
+        logger.info("Starting dataset preparation with validation...")
+        
         try:
-            logger.info("Starting dataset preparation process...")
-            
-            datasets = []
-            
-            # Load Kaggle dataset
-            kaggle_df = self.load_kaggle_dataset()
-            if kaggle_df is not None:
-                datasets.append(kaggle_df)
+            # Load and validate datasets
+            kaggle_result = self.load_kaggle_dataset()
+            liar_result = self.load_liar_dataset()
             
-            # Load LIAR dataset
-            liar_df = self.load_liar_dataset()
-            if liar_df is not None:
-                datasets.append(liar_df)
-            
-            # Check if we have any data
-            if not datasets:
-                error_msg = "No datasets could be loaded successfully"
-                logger.error(error_msg)
-                return False, error_msg
+            # Handle None returns gracefully
+            if kaggle_result is None:
+                logger.warning("Kaggle dataset loading failed")
+                kaggle_df, kaggle_validation = pd.DataFrame(), {
+                    'source': 'kaggle_combined', 'original_count': 0, 'valid_count': 0, 
+                    'success_rate': 0, 'overall_quality_score': 0, 'validation_timestamp': datetime.now().isoformat()
+                }
+            else:
+                kaggle_df, kaggle_validation = kaggle_result
             
-            # Combine all datasets
-            logger.info("Combining all datasets...")
-            combined_df = pd.concat(datasets, ignore_index=True)
+            if liar_result is None:
+                logger.warning("LIAR dataset loading failed")
+                liar_df, liar_validation = pd.DataFrame(), {
+                    'source': 'liar_combined', 'original_count': 0, 'valid_count': 0, 
+                    'success_rate': 0, 'overall_quality_score': 0, 'validation_timestamp': datetime.now().isoformat()
+                }
+            else:
+                liar_df, liar_validation = liar_result
             
-            # Validate and clean the combined dataset
-            validated_df = self.validate_dataset(combined_df)
+            # Combine datasets
+            datasets_to_combine = [df for df in [kaggle_df, liar_df] if not df.empty]
             
-            # Check minimum sample requirement
-            if len(validated_df) < 100:
-                error_msg = f"Insufficient samples after validation: {len(validated_df)}"
-                logger.error(error_msg)
-                return False, error_msg
+            if not datasets_to_combine:
+                return False, "No datasets could be loaded and validated"
             
-            # Generate metadata
-            metadata = self.generate_dataset_metadata(validated_df)
+            combined_df = pd.concat(datasets_to_combine, ignore_index=True)
             
-            # Save dataset
-            validated_df.to_csv(self.output_path, index=False)
+            # Save combined dataset
+            combined_df.to_csv(self.output_path, index=False)
             
-            # Save metadata
-            with open(self.metadata_path, 'w') as f:
-                json.dump(metadata, f, indent=2)
+            # Save validation reports
+            total_original = kaggle_validation['original_count'] + liar_validation['original_count']
+            validation_report = {
+                'datasets': {
+                    'kaggle': kaggle_validation,
+                    'liar': liar_validation
+                },
+                'combined_stats': {
+                    'total_articles': len(combined_df),
+                    'total_original': total_original,
+                    'overall_success_rate': len(combined_df) / max(1, total_original),
+                    'validation_timestamp': datetime.now().isoformat()
+                }
+            }
             
-            logger.info(f"Dataset preparation complete!")
-            logger.info(f"Final dataset: {len(validated_df)} samples")
-            logger.info(f"Quality score: {metadata['quality_score']:.3f}")
-            logger.info(f"Saved to: {self.output_path}")
+            validation_report_path = self.output_dir / "dataset_validation_report.json"
+            with open(validation_report_path, 'w') as f:
+                json.dump(validation_report, f, indent=2)
             
-            return True, f"Successfully prepared {len(validated_df)} samples"
+            logger.info(f"Dataset preparation complete. Validation report saved to {validation_report_path}")
+            return True, f"Successfully prepared {len(combined_df)} validated articles"
             
         except Exception as e:
-            error_msg = f"Dataset preparation failed: {str(e)}"
-            logger.error(error_msg)
-            return False, error_msg
+            logger.error(f"Dataset preparation failed: {e}")
+            return False, f"Dataset preparation failed: {str(e)}"
 
+    def validate_dataset_with_schemas(self, df: pd.DataFrame, source_name: str) -> Tuple[pd.DataFrame, Dict]:
+        """Validate dataset using comprehensive schemas"""
+        logger.info(f"Starting schema validation for {source_name}...")
+        
+        validator = DataValidationPipeline()
+        
+        # Convert DataFrame to validation format
+        articles_data = []
+        for _, row in df.iterrows():
+            article_data = {
+                'text': str(row.get('text', '')),
+                'label': int(row.get('label', 0)),
+                'source': source_name
+            }
+            
+            if 'title' in row and pd.notna(row['title']):
+                article_data['title'] = str(row['title'])
+            if 'url' in row and pd.notna(row['url']):
+                article_data['url'] = str(row['url'])
+                
+            articles_data.append(article_data)
+        
+        # Perform batch validation
+        validation_result = validator.validate_batch(
+            articles_data, 
+            batch_id=f"{source_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
+            validation_level=ValidationLevel.MODERATE
+        )
+        
+        # Filter valid articles and add quality scores
+        valid_indices = [i for i, result in enumerate(validation_result.validation_results) if result.is_valid]
+        
+        if valid_indices:
+            valid_df = df.iloc[valid_indices].copy()
+            quality_scores = [validation_result.validation_results[i].quality_metrics.get('overall_quality_score', 0.0) 
+                             for i in valid_indices]
+            valid_df['validation_quality_score'] = quality_scores
+            valid_df['validation_timestamp'] = datetime.now().isoformat()
+        else:
+            valid_df = pd.DataFrame(columns=df.columns)
+        
+        validation_summary = {
+            'source': source_name,
+            'original_count': len(df),
+            'valid_count': len(valid_df),
+            'success_rate': validation_result.success_rate,
+            'overall_quality_score': validation_result.overall_quality_score,
+            'validation_timestamp': datetime.now().isoformat()
+        }
+        
+        return valid_df, validation_summary
+    
+    
+    
 def main():
     """Main execution function"""
     preparer = DatasetPreparer()