Update app.py

app.py

@@ -5,20 +5,166 @@ import os
 import tempfile
 import time
 import subprocess
+import json
 from huggingface_hub import login, HfApi
 from transformers import AutoTokenizer, BertForSequenceClassification
-from datasets import load_dataset
+from datasets import load_dataset, Dataset, DatasetDict
 
 # Global variables
 MODEL_PATH = "local-model"
 CATEGORIES = ['Online-Safety', 'BroadBand', 'TV-Radio']
 idx_to_category = {0: 'Online-Safety', 1: 'BroadBand', 2: 'TV-Radio'}
+category_to_idx = {'Online-Safety': 0, 'BroadBand': 1, 'TV-Radio': 2}
 TOKEN = None
 TRAINING_LOGS = []
 CURRENT_MODEL = None
 CURRENT_TOKENIZER = None
 
-def login_to_hf(token):
+# Local data files
+LOCAL_DATA_FILES = [
+    "merged-test-data.csv",
+    "test-category.csv", 
+    "test-complaint.csv"
+]
+
+def get_available_datasets():
+    """Get list of available local datasets"""
+    available_files = []
+    for file in LOCAL_DATA_FILES:
+        if os.path.exists(file):
+            try:
+                df = pd.read_csv(file)
+                available_files.append(f"{file} ({len(df)} rows)")
+            except Exception as e:
+                available_files.append(f"{file} (Error: {str(e)})")
+        else:
+            available_files.append(f"{file} (Not found)")
+    
+    # Also check for any other CSV files in the directory
+    for file in os.listdir("."):
+        if file.endswith(".csv") and file not in LOCAL_DATA_FILES:
+            if os.path.exists(file):
+                try:
+                    df = pd.read_csv(file)
+                    available_files.append(f"{file} ({len(df)} rows)")
+                except:
+                    available_files.append(f"{file} (Error reading)")
+    
+    return available_files
+
+def load_and_prepare_local_dataset(file_path, text_column, label_column, test_size=0.2):
+    """Load and prepare local CSV dataset for training"""
+    try:
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Dataset file not found: {file_path}")
+        
+        # Load the CSV file
+        df = pd.read_csv(file_path)
+        
+        # Verify required columns exist
+        if text_column not in df.columns:
+            available_cols = list(df.columns)
+            raise ValueError(f"Text column '{text_column}' not found. Available columns: {available_cols}")
+        
+        if label_column not in df.columns:
+            available_cols = list(df.columns)
+            raise ValueError(f"Label column '{label_column}' not found. Available columns: {available_cols}")
+        
+        # Clean the data
+        df = df.dropna(subset=[text_column, label_column])
+        df[text_column] = df[text_column].astype(str)
+        
+        # Handle different label formats
+        if df[label_column].dtype == 'object':
+            # If labels are text, convert to indices
+            unique_labels = df[label_column].unique()
+            if len(unique_labels) > len(CATEGORIES):
+                raise ValueError(f"Too many unique labels ({len(unique_labels)}). Expected max {len(CATEGORIES)}")
+            
+            # Try to map text labels to our categories
+            label_mapping = {}
+            for label in unique_labels:
+                if label in category_to_idx:
+                    label_mapping[label] = category_to_idx[label]
+                else:
+                    # Auto-assign if not found
+                    available_indices = set(range(len(CATEGORIES))) - set(label_mapping.values())
+                    if available_indices:
+                        label_mapping[label] = min(available_indices)
+                    else:
+                        raise ValueError(f"Cannot map label '{label}' to available categories")
+            
+            df['label_idx'] = df[label_column].map(label_mapping)
+        else:
+            # If labels are already numeric
+            df['label_idx'] = df[label_column].astype(int)
+        
+        # Verify label indices are valid
+        if df['label_idx'].min() < 0 or df['label_idx'].max() >= len(CATEGORIES):
+            raise ValueError(f"Label indices must be between 0 and {len(CATEGORIES)-1}")
+        
+        # Create train/validation split
+        from sklearn.model_selection import train_test_split
+        
+        train_df, val_df = train_test_split(
+            df, 
+            test_size=test_size, 
+            random_state=42, 
+            stratify=df['label_idx']
+        )
+        
+        # Convert to Hugging Face datasets
+        train_dataset = Dataset.from_pandas(train_df[[text_column, 'label_idx']])
+        val_dataset = Dataset.from_pandas(val_df[[text_column, 'label_idx']])
+        
+        dataset_dict = DatasetDict({
+            'train': train_dataset,
+            'validation': val_dataset
+        })
+        
+        return dataset_dict, text_column, 'label_idx'
+        
+    except Exception as e:
+        raise Exception(f"Error loading dataset: {str(e)}")
+
+def preview_dataset(file_path, text_column, label_column):
+    """Preview a dataset file"""
+    try:
+        if not file_path:
+            return "Please select a dataset file first."
+        
+        if not os.path.exists(file_path):
+            return f"❌ File not found: {file_path}"
+        
+        df = pd.read_csv(file_path)
+        
+        preview_info = []
+        preview_info.append(f"📊 **Dataset Preview: {file_path}**")
+        preview_info.append(f"- **Total rows:** {len(df)}")
+        preview_info.append(f"- **Columns:** {list(df.columns)}")
+        preview_info.append("")
+        
+        if text_column in df.columns:
+            preview_info.append(f"✅ **Text column '{text_column}' found**")
+            preview_info.append(f"- Sample text: {str(df[text_column].iloc[0])[:100]}...")
+        else:
+            preview_info.append(f"❌ **Text column '{text_column}' not found**")
+            return "\n".join(preview_info)
+        
+        if label_column in df.columns:
+            preview_info.append(f"✅ **Label column '{label_column}' found**")
+            label_counts = df[label_column].value_counts()
+            preview_info.append("- **Label distribution:**")
+            for label, count in label_counts.items():
+                preview_info.append(f"  - {label}: {count} ({count/len(df)*100:.1f}%)")
+        else:
+            preview_info.append(f"❌ **Label column '{label_column}' not found**")
+            return "\n".join(preview_info)
+        
+        return "\n".join(preview_info)
+        
+    except Exception as e:
+        return f"❌ Error previewing dataset: {str(e)}"
     """Login to Hugging Face"""
     global TOKEN
     TOKEN = token
@@ -162,9 +308,9 @@ def predict_csv(csv_file, model_path):
     except Exception as e:
         return f"❌ CSV processing failed: {str(e)}"
 
-def train_model(dataset_name, num_epochs, batch_size, learning_rate, hf_token, 
-                push_to_hub, username, model_name):
-    """Start the model training process"""
+def train_model(dataset_file, text_column, label_column, num_epochs, batch_size, 
+                learning_rate, hf_token, push_to_hub, username, model_name):
+    """Start the model training process with local data"""
     global TRAINING_LOGS, MODEL_PATH
     
     TRAINING_LOGS = []  # Reset logs at the start of training
@@ -184,14 +330,47 @@ def train_model(dataset_name, num_epochs, batch_size, learning_rate, hf_token,
     else:
         hub_model_id = None
     
-    # Create training command
+    # Validate dataset file
+    if not dataset_file or not os.path.exists(dataset_file):
+        TRAINING_LOGS.append(f"❌ Dataset file not found: {dataset_file}")
+        yield "\n".join(TRAINING_LOGS)
+        return
+    
+    try:
+        # Load and prepare the dataset
+        TRAINING_LOGS.append(f"📊 Loading dataset from {dataset_file}...")
+        yield "\n".join(TRAINING_LOGS)
+        
+        dataset_dict, final_text_col, final_label_col = load_and_prepare_local_dataset(
+            dataset_file, text_column, label_column
+        )
+        
+        TRAINING_LOGS.append(f"✅ Dataset loaded successfully!")
+        TRAINING_LOGS.append(f"- Train samples: {len(dataset_dict['train'])}")
+        TRAINING_LOGS.append(f"- Validation samples: {len(dataset_dict['validation'])}")
+        yield "\n".join(TRAINING_LOGS)
+        
+        # Save dataset temporarily for the training script
+        temp_dataset_path = "temp_dataset"
+        os.makedirs(temp_dataset_path, exist_ok=True)
+        dataset_dict.save_to_disk(temp_dataset_path)
+        
+        TRAINING_LOGS.append("💾 Dataset prepared for training...")
+        yield "\n".join(TRAINING_LOGS)
+        
+    except Exception as e:
+        TRAINING_LOGS.append(f"❌ Error preparing dataset: {str(e)}")
+        yield "\n".join(TRAINING_LOGS)
+        return
+    
+    # Create training command for local dataset
     cmd = [
         "python", "bert_finetune.py",
-        "--dataset_name", dataset_name,
+        "--dataset_path", temp_dataset_path,  # Use local path instead of HF dataset name
         "--model_id", "bert-base-uncased",
         "--output_dir", MODEL_PATH,
-        "--feature_column", "complaint",
-        "--label_column", "label_idx",
+        "--feature_column", final_text_col,
+        "--label_column", final_label_col,
         "--num_labels", "3",
         "--num_train_epochs", str(num_epochs),
         "--batch_size", str(batch_size),
@@ -204,7 +383,7 @@ def train_model(dataset_name, num_epochs, batch_size, learning_rate, hf_token,
         if hf_token:
             cmd.extend(["--hf_token", hf_token])
     
-    TRAINING_LOGS.append(f"Starting training with command: {' '.join(cmd)}")
+    TRAINING_LOGS.append(f"🚀 Starting training with command: {' '.join(cmd)}")
     yield "\n".join(TRAINING_LOGS)
     
     try:
@@ -216,7 +395,7 @@ def train_model(dataset_name, num_epochs, batch_size, learning_rate, hf_token,
             bufsize=1
         )
         
-        TRAINING_LOGS.append("Training started...")
+        TRAINING_LOGS.append("🔄 Training started...")
         yield "\n".join(TRAINING_LOGS)
         
         while True:
@@ -232,13 +411,21 @@ def train_model(dataset_name, num_epochs, batch_size, learning_rate, hf_token,
         if process.returncode == 0:
             TRAINING_LOGS.append("✅ Training completed successfully!")
             if push_to_hub and hub_model_id:
-                TRAINING_LOGS.append(f"✅ Model pushed to Hugging Face Hub: {hub_model_id}")
+                TRAINING_LOGS.append(f"🤗 Model pushed to Hugging Face Hub: {hub_model_id}")
             
             # Load the trained model
-            TRAINING_LOGS.append("Loading trained model...")
+            TRAINING_LOGS.append("📥 Loading trained model...")
             load_result = load_model(MODEL_PATH)
             TRAINING_LOGS.append(load_result)
             
+            # Clean up temporary files
+            import shutil
+            try:
+                shutil.rmtree(temp_dataset_path)
+                TRAINING_LOGS.append("🧹 Cleaned up temporary files")
+            except:
+                pass
+            
             # Final success message
             TRAINING_LOGS.append("\n✨ All done! Your model is ready to use.")
         else:
@@ -289,14 +476,40 @@ with gr.Blocks(title="BERT Complaint Classifier") as app:
     with gr.Tabs():
         # Training Tab
         with gr.TabItem("Train Model"):
-            gr.Markdown("### Train a New Model")
-            gr.Markdown("Provide your dataset information and training parameters")
+            gr.Markdown("### Train a New Model with Local Data")
+            gr.Markdown("Select your local CSV file and configure training parameters")
             
-            dataset_name = gr.Textbox(
-                label="Dataset Name (from Hugging Face)", 
-                placeholder="e.g., your-username/complaint-categories-dataset"
-            )
+            # Dataset selection and preview
+            with gr.Row():
+                with gr.Column(scale=2):
+                    dataset_file = gr.Dropdown(
+                        label="Select Dataset File",
+                        choices=[f for f in os.listdir(".") if f.endswith(".csv")],
+                        value=LOCAL_DATA_FILES[0] if LOCAL_DATA_FILES[0] in os.listdir(".") else None,
+                        allow_custom_value=True
+                    )
+                    
+                with gr.Column(scale=1):
+                    refresh_btn = gr.Button("🔄 Refresh Files", size="sm")
+                    
+            # Column configuration
+            with gr.Row():
+                text_column = gr.Textbox(
+                    label="Text Column Name", 
+                    value="complaint",
+                    placeholder="e.g., complaint, text, description"
+                )
+                label_column = gr.Textbox(
+                    label="Label Column Name", 
+                    value="category",
+                    placeholder="e.g., category, label, class"
+                )
+            
+            # Dataset preview
+            preview_btn = gr.Button("📊 Preview Dataset", variant="secondary")
+            dataset_preview = gr.Markdown("Select a dataset file and click 'Preview Dataset' to see its structure.")
             
+            # Training parameters
             with gr.Row():
                 num_epochs = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Number of Epochs")
                 batch_size = gr.Slider(minimum=4, maximum=32, value=8, step=4, label="Batch Size")