Create clean_dataset.py

clean_dataset.py

@@ -0,0 +1,447 @@
+import pandas as pd
+import numpy as np
+import json
+import os
+import argparse
+import logging
+from tqdm import tqdm
+import chardet
+import csv
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("dataset_cleaner.log"),
+        logging.StreamHandler()
+    ]
+)
+
+logger = logging.getLogger(__name__)
+
+class SaaSDatasetCleaner:
+    """
+    Class for cleaning and validating the SaaS sales conversation dataset.
+    Handles issues resulting from interrupted generations.
+    """
+    
+    def __init__(self, input_file, output_file=None, chunk_size=1000, encoding='utf-8', skip_encoding_check=False):
+        """
+        Initialize the cleaner.
+        
+        Args:
+            input_file: Path to the input CSV file
+            output_file: Path to save cleaned dataset (defaults to 'cleaned_' + input_file)
+            chunk_size: Number of rows to process at once
+            encoding: File encoding (defaults to utf-8)
+            skip_encoding_check: Whether to skip encoding detection and line-by-line processing
+        """
+        self.input_file = input_file
+        self.output_file = output_file or f"cleaned_{os.path.basename(input_file)}"
+        self.chunk_size = chunk_size
+        self.encoding = encoding
+        self.skip_encoding_check = skip_encoding_check
+        self.stats = {
+            'total_rows': 0,
+            'valid_rows': 0,
+            'invalid_json': 0,
+            'missing_values': 0,
+            'invalid_embeddings': 0,
+            'duplicates': 0,
+            'encoding_errors': 0,
+            'recovered_rows': 0
+        }
+        
+        # If not skipping encoding check, detect encoding
+        if not self.skip_encoding_check and not self.encoding:
+            self.detect_encoding()
+        
+        # Get the columns and prepare for processing
+        self.initialize_columns()
+    
+    def detect_encoding(self):
+        """Detect the file encoding."""
+        logger.info("Detecting file encoding...")
+        
+        # Read a sample of the file to detect encoding
+        with open(self.input_file, 'rb') as f:
+            sample = f.read(min(10000000, os.path.getsize(self.input_file)))  # Read up to 10MB
+        
+        result = chardet.detect(sample)
+        self.encoding = result['encoding']
+        confidence = result['confidence']
+        
+        logger.info(f"Detected encoding: {self.encoding} with confidence: {confidence:.2f}")
+        
+        # If confidence is low, try common encodings
+        if confidence < 0.7:
+            logger.warning(f"Low confidence in encoding detection. Will try multiple encodings.")
+            self.encoding = None  # Will try multiple encodings later
+    
+    def initialize_columns(self):
+        """Initialize column information."""
+        # Try to read the header with different encodings if needed
+        encodings_to_try = ['utf-8'] if (self.skip_encoding_check or self.encoding) else ['utf-8', 'latin1', 'iso-8859-1', 'cp1252']
+        
+        for enc in encodings_to_try:
+            try:
+                # Try to read just the header
+                with open(self.input_file, 'r', encoding=enc, errors='replace') as f:
+                    reader = csv.reader(f)
+                    self.columns = next(reader)
+                
+                self.encoding = enc
+                logger.info(f"Successfully read header with encoding: {enc}")
+                
+                # Identify embedding columns
+                self.embedding_cols = [col for col in self.columns if col.startswith('embedding_')]
+                logger.info(f"Found {len(self.embedding_cols)} embedding columns")
+                
+                return
+                
+            except Exception as e:
+                logger.warning(f"Failed to read header with encoding {enc}: {str(e)}")
+        
+        # If we get here, all encodings failed
+        logger.error("Could not read column headers with any encoding")
+        self.columns = []
+        self.embedding_cols = []
+    
+    def process_line_by_line(self):
+        """Process the file line by line to handle encoding issues."""
+        logger.info("Processing file line by line to handle encoding issues...")
+        
+        # Open the output file
+        with open(self.output_file, 'w', encoding='utf-8', newline='') as out_file:
+            writer = None  # Will initialize after getting headers
+            
+            # Process the input file
+            with open(self.input_file, 'rb') as in_file:
+                # Process line by line
+                line_count = 0
+                valid_count = 0
+                
+                for line in tqdm(in_file, desc="Reading lines"):
+                    line_count += 1
+                    
+                    # Try to decode with multiple encodings
+                    decoded_line = None
+                    for enc in ['utf-8', 'latin1', 'iso-8859-1', 'cp1252']:
+                        try:
+                            decoded_line = line.decode(enc)
+                            break
+                        except UnicodeDecodeError:
+                            continue
+                    
+                    if decoded_line is None:
+                        # Could not decode with any encoding, skip line
+                        self.stats['encoding_errors'] += 1
+                        continue
+                    
+                    # Parse the CSV line
+                    try:
+                        reader = csv.reader([decoded_line])
+                        row = next(reader)
+                        
+                        # Initialize writer with headers if this is the first line
+                        if line_count == 1:
+                            writer = csv.writer(out_file)
+                            writer.writerow(row)  # Write headers
+                            continue
+                        
+                        # Basic validation - check number of columns
+                        if len(row) != len(self.columns):
+                            logger.debug(f"Line {line_count}: Column count mismatch. Expected {len(self.columns)}, got {len(row)}")
+                            continue
+                        
+                        # Write the row
+                        writer.writerow(row)
+                        valid_count += 1
+                        
+                    except Exception as e:
+                        logger.debug(f"Error processing line {line_count}: {str(e)}")
+                        self.stats['encoding_errors'] += 1
+            
+            self.stats['total_rows'] = line_count - 1  # Subtract header
+            self.stats['recovered_rows'] = valid_count
+            
+            logger.info(f"Processed {line_count} lines, recovered {valid_count} valid rows")
+            logger.info(f"Found {self.stats['encoding_errors']} lines with encoding errors")
+    
+    def _validate_json_fields(self, df):
+        """Validate and clean JSON fields."""
+        # List of columns that should contain JSON
+        json_columns = ['scenario', 'conversation', 'probability_trajectory']
+        
+        for col in json_columns:
+            if col not in df.columns:
+                continue
+                
+            # Create a valid indicator
+            df[f'{col}_valid'] = True
+            
+            # Check each value
+            for idx, value in enumerate(df[col]):
+                try:
+                    if pd.isna(value):
+                        df.at[idx, f'{col}_valid'] = False
+                        self.stats['invalid_json'] += 1
+                        continue
+                    
+                    # Attempt to parse JSON
+                    json.loads(value)
+                except:
+                    df.at[idx, f'{col}_valid'] = False
+                    self.stats['invalid_json'] += 1
+        
+        # Create an overall valid flag
+        valid_flags = [f'{col}_valid' for col in json_columns if f'{col}_valid' in df.columns]
+        if valid_flags:
+            df['json_valid'] = df[valid_flags].all(axis=1)
+        else:
+            df['json_valid'] = True
+            
+        # Clean up the temporary columns
+        for col in json_columns:
+            if f'{col}_valid' in df.columns:
+                df = df.drop(columns=[f'{col}_valid'])
+                
+        return df
+    
+    def _validate_embeddings(self, df):
+        """Check if embeddings are valid."""
+        if not self.embedding_cols:
+            return df
+        
+        # Check if the first embedding column has a value as a simple check
+        if 'embedding_0' in df.columns:
+            df['embeddings_valid'] = ~df['embedding_0'].isna()
+        else:
+            df['embeddings_valid'] = True
+            
+        # Count invalid embeddings
+        self.stats['invalid_embeddings'] += (~df['embeddings_valid']).sum()
+            
+        return df
+    
+    def _check_missing_values(self, df):
+        """Check for missing values in important columns."""
+        important_cols = [
+            'company_id', 'company_name', 'product_name', 'conversation_id', 
+            'conversation', 'full_text', 'outcome'
+        ]
+        
+        # Filter to columns that actually exist
+        important_cols = [col for col in important_cols if col in df.columns]
+        
+        if not important_cols:
+            df['missing_important'] = False
+            return df
+        
+        # Create a flag for rows with missing important values
+        missing_flags = df[important_cols].isna().any(axis=1)
+        df['missing_important'] = missing_flags
+        
+        # Count missing values
+        self.stats['missing_values'] += missing_flags.sum()
+        
+        return df
+    
+    def _flag_valid_rows(self, df):
+        """Create a single flag for valid rows."""
+        # A row is valid if it has valid JSON, valid embeddings, and no missing important values
+        required_flags = []
+        
+        if 'json_valid' in df.columns:
+            required_flags.append('json_valid')
+            
+        if 'embeddings_valid' in df.columns:
+            required_flags.append('embeddings_valid')
+            
+        if 'missing_important' in df.columns:
+            required_flags.append('~missing_important')
+        
+        if required_flags:
+            if '~missing_important' in required_flags:
+                required_flags.remove('~missing_important')
+                if required_flags:
+                    df['row_valid'] = df[required_flags].all(axis=1) & ~df['missing_important']
+                else:
+                    df['row_valid'] = ~df['missing_important']
+            else:
+                df['row_valid'] = df[required_flags].all(axis=1)
+        else:
+            df['row_valid'] = True
+        
+        # Update valid rows count
+        self.stats['valid_rows'] += df['row_valid'].sum()
+        
+        return df
+    
+    def _remove_duplicates(self, df):
+        """Remove duplicate conversation IDs."""
+        if 'conversation_id' in df.columns:
+            # Check for duplicates
+            dup_mask = df.duplicated(subset=['conversation_id'], keep='first')
+            df['is_duplicate'] = dup_mask
+            
+            # Count duplicates
+            self.stats['duplicates'] += dup_mask.sum()
+        else:
+            df['is_duplicate'] = False
+            
+        return df
+    
+    def clean_dataset(self):
+        """
+        Clean the dataset by first fixing encoding issues, then cleaning the data.
+        """
+        logger.info(f"Starting to clean dataset: {self.input_file}")
+        
+        # Check if the file exists
+        if not os.path.exists(self.input_file):
+            logger.error(f"Input file not found: {self.input_file}")
+            return
+        
+        # If we're not skipping encoding checks, process line by line
+        if not self.skip_encoding_check:
+            self.process_line_by_line()
+            intermediate_file = self.output_file
+            self.output_file = f"validated_{os.path.basename(self.input_file)}"
+        else:
+            logger.info("Skipping encoding check as requested")
+            # Use the input file directly as the intermediate file
+            intermediate_file = self.input_file
+            
+            # Count rows in the file for progress tracking
+            with open(intermediate_file, 'r', encoding=self.encoding) as f:
+                self.stats['total_rows'] = sum(1 for _ in f) - 1  # Subtract header
+                self.stats['recovered_rows'] = self.stats['total_rows']
+                
+            logger.info(f"Total rows to validate: {self.stats['total_rows']}")
+        
+        # Now that we have a cleaned file with proper encoding, process it for data validation
+        logger.info("Beginning data validation on recovered rows...")
+        
+        # Get the total number of rows for progress tracking
+        try:
+            total_rows = self.stats['recovered_rows']
+            logger.info(f"Total rows to validate: {total_rows}")
+        except Exception as e:
+            logger.error(f"Error counting rows: {str(e)}")
+            total_rows = 0
+        
+        # Process the dataset in chunks
+        try:
+            # Create a reader - now with known proper encoding
+            # Use error_bad_lines=False for older pandas versions (renamed to on_bad_lines in newer versions)
+            reader = pd.read_csv(
+                intermediate_file, 
+                chunksize=self.chunk_size,
+                encoding='utf-8',
+                low_memory=False,  # Avoid dtype warnings
+                error_bad_lines=False  # Skip bad lines (older parameter name)
+            )
+            
+            # Create a header flag for the first chunk
+            first_chunk = True
+            
+            # Process each chunk
+            with tqdm(total=total_rows, desc="Validating data") as pbar:
+                for chunk_num, chunk in enumerate(reader):
+                    logger.debug(f"Processing chunk {chunk_num+1}")
+                    
+                    # Run validation steps
+                    chunk = self._validate_json_fields(chunk)
+                    chunk = self._validate_embeddings(chunk)
+                    chunk = self._check_missing_values(chunk)
+                    chunk = self._remove_duplicates(chunk)
+                    chunk = self._flag_valid_rows(chunk)
+                    
+                    # Filter to valid rows only
+                    valid_chunk = chunk[chunk['row_valid'] & ~chunk['is_duplicate']]
+                    
+                    # Remove the validation columns
+                    for col in ['json_valid', 'embeddings_valid', 'missing_important', 'row_valid', 'is_duplicate']:
+                        if col in valid_chunk.columns:
+                            valid_chunk = valid_chunk.drop(columns=[col])
+                    
+                    # Write the cleaned chunk
+                    valid_chunk.to_csv(
+                        self.output_file,
+                        mode='w' if first_chunk else 'a',
+                        header=first_chunk,
+                        index=False,
+                        encoding='utf-8'
+                    )
+                    
+                    # Update the first chunk flag
+                    if first_chunk:
+                        first_chunk = False
+                    
+                    # Update progress
+                    pbar.update(len(chunk))
+            
+            logger.info(f"Dataset cleaning complete. Results saved to {self.output_file}")
+            
+            # Print statistics
+            logger.info(f"Cleaning Statistics:")
+            logger.info(f"- Total rows processed: {self.stats['total_rows']}")
+            logger.info(f"- Rows recovered from encoding issues: {self.stats['recovered_rows']}")
+            logger.info(f"- Encoding errors: {self.stats['encoding_errors']}")
+            logger.info(f"- Valid rows after validation: {self.stats['valid_rows']}")
+            logger.info(f"- Rows with invalid JSON: {self.stats['invalid_json']}")
+            logger.info(f"- Rows with missing values: {self.stats['missing_values']}")
+            logger.info(f"- Rows with invalid embeddings: {self.stats['invalid_embeddings']}")
+            logger.info(f"- Duplicate rows: {self.stats['duplicates']}")
+            
+            # Create a summary file
+            with open(f"{self.output_file}_summary.txt", 'w') as f:
+                f.write("Dataset Cleaning Summary\n")
+                f.write("=======================\n\n")
+                f.write(f"Input file: {self.input_file}\n")
+                f.write(f"Output file: {self.output_file}\n\n")
+                f.write(f"Total rows processed: {self.stats['total_rows']}\n")
+                f.write(f"Rows recovered from encoding issues: {self.stats['recovered_rows']}\n")
+                f.write(f"Encoding errors: {self.stats['encoding_errors']}\n")
+                f.write(f"Valid rows after validation: {self.stats['valid_rows']}\n")
+                f.write(f"Rows with invalid JSON: {self.stats['invalid_json']}\n")
+                f.write(f"Rows with missing values: {self.stats['missing_values']}\n")
+                f.write(f"Rows with invalid embeddings: {self.stats['invalid_embeddings']}\n")
+                f.write(f"Duplicate rows: {self.stats['duplicates']}\n")
+            
+            return self.stats
+            
+        except Exception as e:
+            logger.error(f"Error validating dataset: {str(e)}")
+            raise e
+
+def main():
+    """Main function to run the dataset cleaner."""
+    parser = argparse.ArgumentParser(description="Clean and validate SaaS sales conversation dataset")
+    parser.add_argument("input_file", type=str, help="Path to the input CSV file")
+    parser.add_argument("--output_file", type=str, default=None, 
+                       help="Path to save cleaned dataset (defaults to 'cleaned_' + input_file)")
+    parser.add_argument("--chunk_size", type=int, default=1000,
+                       help="Number of rows to process at once")
+    parser.add_argument("--encoding", type=str, default='utf-8',
+                       help="File encoding (defaults to utf-8)")
+    parser.add_argument("--skip_encoding_check", action="store_true",
+                       help="Skip encoding detection and line-by-line processing")
+    
+    args = parser.parse_args()
+    
+    # Create and run the cleaner
+    cleaner = SaaSDatasetCleaner(
+        input_file=args.input_file,
+        output_file=args.output_file,
+        chunk_size=args.chunk_size,
+        encoding=args.encoding,
+        skip_encoding_check=args.skip_encoding_check
+    )
+    
+    cleaner.clean_dataset()
+
+if __name__ == "__main__":
+    main()