Add training app with Flan-T5 implementation and datasets

README.md

@@ -1,13 +1,48 @@
 ---
-title: Win Stack
-emoji: 🔥
-colorFrom: green
-colorTo: indigo
+title: Resume Normalizer Trainer
+emoji: 📝
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 5.37.0
+sdk_version: 5.6.0
 app_file: app.py
 pinned: false
-short_description: Winstack
+license: apache-2.0
+hardware: 4xL4
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Resume Normalizer Trainer
+
+Fine-tune a Flan-T5 model for resume entity normalization and deduplication.
+
+## Features
+
+- **Company Name Normalization**: Handle mergers, acquisitions, and rebranding (e.g., "Facebook" → "Meta Platforms Inc.")
+- **Job Title Standardization**: Recognize equivalent roles and seniority levels (e.g., "SWE" → "Software Engineer")
+- **Skills Normalization**: Standardize technology names and abbreviations (e.g., "JS" → "JavaScript")
+- **Binary Equivalency Detection**: Determine if two entities refer to the same thing
+
+## Model Details
+
+- **Base Model**: Google Flan-T5 (instruction-tuned for better zero-shot performance)
+- **Fine-tuning Method**: LoRA (Low-Rank Adaptation) for efficient training
+- **Parameters**: 250M (T5-Base) or 770M (T5-Large)
+- **Training Data**: 9,302 high-quality examples (478 manual + 8,824 synthetic)
+
+## Usage
+
+1. Check that training data is available using the "Check Data" tab
+2. Enter your HuggingFace token and username
+3. Select model size and training epochs
+4. Click "Start Training" and monitor progress in the "Training Status" tab
+5. Once complete, your model will be available on HuggingFace Hub
+
+## Expected Performance
+
+- **Inference Speed**: <100ms per query
+- **Accuracy**: >90% on entity normalization tasks
+- **Memory Usage**: ~1GB (T5-Base) or ~3GB (T5-Large)
+
+## Hardware Requirements
+
+This Space runs on 4xL4 GPUs (96GB total VRAM) for efficient distributed training.

app.py

@@ -1,7 +1,252 @@
 import gradio as gr
+import os
+import subprocess
+import threading
+import time
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Global variable to track training status
+training_status = {"status": "idle", "message": "", "progress": 0}
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+def check_data():
+    """Check if data is available"""
+    files = []
+    if os.path.exists("combined_final_training_data.csv"):
+        files.append("✅ Combined dataset: 9,302 examples")
+    if os.path.exists("combined_balanced_training_data.csv"):
+        files.append("✅ Balanced dataset: 8,304 examples")
+    if os.path.exists("data/clean_training_data.csv"):
+        files.append("✅ Clean manual data: 478 examples")
+    
+    if not files:
+        return "❌ No training data found. Please upload data files."
+    
+    return "\n".join(files)
+
+def run_training_subprocess(hf_token, model_size, hub_username, num_epochs, use_balanced):
+    """Run training in subprocess"""
+    global training_status
+    
+    try:
+        # Determine which data file to use
+        if use_balanced and os.path.exists("combined_balanced_training_data.csv"):
+            data_path = "combined_balanced_training_data.csv"
+        elif os.path.exists("combined_final_training_data.csv"):
+            data_path = "combined_final_training_data.csv"
+        else:
+            training_status["status"] = "error"
+            training_status["message"] = "No training data found!"
+            return
+        
+        # Determine model size
+        size = "base" if "Base" in model_size else "large"
+        
+        # Build command
+        cmd = [
+            "python", "train.py",
+            "--data_path", data_path,
+            "--model_size", size,
+            "--num_epochs", str(num_epochs),
+            "--use_lora"  # Always use LoRA for efficiency
+        ]
+        
+        if hf_token:
+            cmd.extend(["--hf_token", hf_token])
+        if hub_username:
+            cmd.extend(["--hub_username", hub_username])
+        
+        training_status["status"] = "running"
+        training_status["message"] = "Starting training..."
+        training_status["progress"] = 0
+        
+        # Run training
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1
+        )
+        
+        # Read output line by line
+        for line in process.stdout:
+            if "loss" in line.lower():
+                training_status["message"] = line.strip()
+            elif "epoch" in line.lower():
+                # Try to extract progress
+                try:
+                    if "/" in line:
+                        parts = line.split("/")
+                        current = float(parts[0].split()[-1])
+                        total = float(parts[1].split()[0])
+                        training_status["progress"] = int((current / total) * 100)
+                except:
+                    pass
+            elif "exact_match" in line.lower():
+                training_status["message"] = f"Evaluation: {line.strip()}"
+        
+        process.wait()
+        
+        if process.returncode == 0:
+            training_status["status"] = "completed"
+            training_status["message"] = "Training completed successfully! Model pushed to HuggingFace Hub."
+            training_status["progress"] = 100
+        else:
+            training_status["status"] = "error"
+            training_status["message"] = f"Training failed with exit code {process.returncode}"
+            
+    except Exception as e:
+        training_status["status"] = "error"
+        training_status["message"] = f"Error: {str(e)}"
+
+def train_model(hf_token, model_size, hub_username, num_epochs, use_balanced):
+    """Start training in background thread"""
+    global training_status
+    
+    if not hf_token:
+        return "❌ Please provide HuggingFace token"
+    
+    if training_status["status"] == "running":
+        return "⚠️ Training already in progress!"
+    
+    # Start training in background thread
+    thread = threading.Thread(
+        target=run_training_subprocess,
+        args=(hf_token, model_size, hub_username, num_epochs, use_balanced)
+    )
+    thread.start()
+    
+    return "🚀 Training started! Check status below..."
+
+def get_training_status():
+    """Get current training status"""
+    global training_status
+    
+    if training_status["status"] == "idle":
+        return "💤 No training in progress"
+    elif training_status["status"] == "running":
+        return f"""
+🏃 Training in progress... ({training_status['progress']}%)
+
+{training_status['message']}
+"""
+    elif training_status["status"] == "completed":
+        return f"""
+✅ Training completed!
+
+{training_status['message']}
+
+Your model is available at: https://huggingface.co/{training_status.get('hub_username', 'your-username')}/resume-normalizer-flan-t5
+"""
+    else:
+        return f"""
+❌ Training failed!
+
+{training_status['message']}
+"""
+
+# Create Gradio interface
+with gr.Blocks(title="Resume Normalizer Trainer") as app:
+    gr.Markdown("""
+    # Resume Normalizer Trainer
+    
+    Train a Flan-T5 model to normalize company names, job titles, and skills from resumes.
+    
+    **Features:**
+    - Company name normalization (e.g., "Google Inc" → "Alphabet Inc.")
+    - Job title standardization (e.g., "SWE" → "Software Engineer")
+    - Skills normalization (e.g., "JS" → "JavaScript")
+    - Binary equivalency detection
+    
+    **Hardware:** Running on 4xL4 GPUs (96GB VRAM)
+    """)
+    
+    with gr.Tab("📊 Check Data"):
+        check_btn = gr.Button("Check Available Datasets", variant="primary")
+        check_output = gr.Textbox(label="Dataset Status", lines=5)
+        check_btn.click(check_data, outputs=check_output)
+    
+    with gr.Tab("🚀 Train Model"):
+        with gr.Row():
+            with gr.Column():
+                hf_token = gr.Textbox(
+                    label="HuggingFace Token",
+                    type="password",
+                    placeholder="hf_...",
+                    info="Required to push model to Hub"
+                )
+                hub_username = gr.Textbox(
+                    label="HuggingFace Username",
+                    value="aoisfhdugbos",
+                    info="Your HuggingFace username"
+                )
+            
+            with gr.Column():
+                model_size = gr.Dropdown(
+                    label="Model Size",
+                    choices=["T5-Base (250M)", "T5-Large (770M)"],
+                    value="T5-Base (250M)",
+                    info="Larger models are more accurate but slower"
+                )
+                num_epochs = gr.Slider(
+                    label="Training Epochs",
+                    minimum=1,
+                    maximum=10,
+                    value=5,
+                    step=1,
+                    info="More epochs = better quality but longer training"
+                )
+        
+        use_balanced = gr.Checkbox(
+            label="Use Balanced Dataset (8,304 examples)",
+            value=False,
+            info="Check to use balanced dataset instead of full dataset (9,302 examples)"
+        )
+        
+        train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg")
+        train_output = gr.Textbox(label="Training Output", lines=5)
+        
+        train_btn.click(
+            train_model,
+            inputs=[hf_token, model_size, hub_username, num_epochs, use_balanced],
+            outputs=train_output
+        )
+    
+    with gr.Tab("📈 Training Status"):
+        status_btn = gr.Button("🔄 Refresh Status", variant="secondary")
+        status_output = gr.Textbox(label="Current Status", lines=10)
+        
+        # Auto-refresh status
+        status_timer = gr.Timer(5)  # Refresh every 5 seconds
+        status_timer.tick(get_training_status, outputs=status_output)
+        
+        status_btn.click(get_training_status, outputs=status_output)
+    
+    with gr.Tab("ℹ️ About"):
+        gr.Markdown("""
+        ## Resume Normalizer Model
+        
+        This trainer fine-tunes a Flan-T5 model for resume entity normalization tasks:
+        
+        ### Supported Tasks:
+        1. **Company Normalization**: Handles mergers, acquisitions, rebranding
+        2. **Job Title Standardization**: Recognizes equivalent roles and seniority
+        3. **Skills Normalization**: Standardizes technology names and abbreviations
+        4. **Equivalency Detection**: Binary classification for entity matching
+        
+        ### Model Architecture:
+        - Base Model: Google Flan-T5 (instruction-tuned)
+        - Fine-tuning: LoRA (Low-Rank Adaptation) for efficiency
+        - Multi-task: Uses task prefixes ([COMPANY], [JOB], [SKILLS])
+        
+        ### Training Data:
+        - 478 manually curated high-quality examples
+        - 8,824 synthetic examples generated with GPT-4
+        - Total: 9,302 training examples
+        
+        ### Expected Performance:
+        - Inference: <100ms per query
+        - Accuracy: >90% on test set
+        - Model size: 250M-770M parameters
+        """)
+
+app.launch()

requirements.txt

@@ -0,0 +1,14 @@
+transformers>=4.36.0
+datasets>=2.16.0
+accelerate>=0.25.0
+peft>=0.7.0
+evaluate>=0.4.1
+rouge_score>=0.1.2
+pandas>=2.0.0
+numpy>=1.24.0
+torch>=2.1.0
+sentencepiece>=0.1.99
+huggingface_hub>=0.20.0
+gradio==5.6.0
+tensorboard>=2.15.0
+scikit-learn>=1.3.0

train.py

@@ -0,0 +1,300 @@
+import torch
+from transformers import (
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+    TrainingArguments,
+    Trainer,
+    DataCollatorForSeq2Seq,
+    EarlyStoppingCallback
+)
+from datasets import Dataset
+import pandas as pd
+import numpy as np
+from accelerate import Accelerator
+import os
+from huggingface_hub import HfFolder
+import logging
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class ResumeNormalizationTrainer:
+    def __init__(self, model_name="google/flan-t5-base", use_lora=True):
+        self.model_name = model_name
+        self.use_lora = use_lora
+        self.accelerator = Accelerator()
+        self.device = self.accelerator.device
+        
+        logger.info(f"Initializing model: {model_name}")
+        logger.info(f"Using LoRA: {use_lora}")
+        logger.info(f"Device: {self.device}")
+        
+        # Load tokenizer and model
+        self.tokenizer = T5Tokenizer.from_pretrained(model_name)
+        self.model = T5ForConditionalGeneration.from_pretrained(model_name)
+        
+        # Setup LoRA if requested
+        if use_lora:
+            self._setup_lora()
+            
+    def _setup_lora(self):
+        """Configure LoRA for efficient fine-tuning"""
+        try:
+            from peft import LoraConfig, get_peft_model, TaskType
+            
+            lora_config = LoraConfig(
+                r=16,  # rank
+                lora_alpha=32,
+                target_modules=["q", "v"],  # T5 attention layers
+                lora_dropout=0.1,
+                bias="none",
+                task_type=TaskType.SEQ_2_SEQ_LM,
+            )
+            
+            self.model = get_peft_model(self.model, lora_config)
+            self.model.print_trainable_parameters()
+            logger.info("LoRA configuration applied successfully")
+        except Exception as e:
+            logger.error(f"Failed to setup LoRA: {e}")
+            raise
+        
+    def load_dataset(self, data_path):
+        """Load and prepare dataset"""
+        logger.info(f"Loading dataset from: {data_path}")
+        df = pd.read_csv(data_path)
+        logger.info(f"Loaded {len(df)} examples")
+        
+        # Add task prefixes if not present
+        def add_task_prefix(row):
+            task = row['task_type']
+            instruction = row['instruction']
+            
+            # Skip if already has prefix
+            if instruction.startswith('['):
+                return instruction
+                
+            if task == 'normalize_company':
+                return f"[COMPANY] {instruction}"
+            elif task == 'normalize_job_title':
+                return f"[JOB] {instruction}"
+            elif task == 'normalize_skill':
+                return f"[SKILLS] {instruction}"
+            elif task == 'company_equivalence':
+                return f"[COMPANY] {instruction}"
+            elif task == 'job_title_equivalence':
+                return f"[JOB] {instruction}"
+            elif task == 'achievement_equivalence':
+                return f"[ACHIEVEMENT] {instruction}"
+            return instruction
+        
+        df['instruction'] = df.apply(add_task_prefix, axis=1)
+        
+        # Split into train/validation
+        train_size = int(0.9 * len(df))
+        train_df = df[:train_size]
+        val_df = df[train_size:]
+        
+        logger.info(f"Train set: {len(train_df)} examples")
+        logger.info(f"Validation set: {len(val_df)} examples")
+        
+        # Convert to HuggingFace Dataset
+        train_dataset = Dataset.from_pandas(train_df)
+        val_dataset = Dataset.from_pandas(val_df)
+        
+        return train_dataset, val_dataset
+    
+    def preprocess_function(self, examples):
+        """Tokenize inputs and targets"""
+        inputs = examples['instruction']
+        targets = examples['output']
+        
+        # Tokenize inputs
+        model_inputs = self.tokenizer(
+            inputs,
+            max_length=256,
+            truncation=True,
+            padding="max_length"
+        )
+        
+        # Tokenize targets
+        with self.tokenizer.as_target_tokenizer():
+            labels = self.tokenizer(
+                targets,
+                max_length=128,
+                truncation=True,
+                padding="max_length"
+            )
+        
+        # Replace padding token id's of the labels by -100
+        labels["input_ids"] = [
+            [(l if l != self.tokenizer.pad_token_id else -100) for l in label]
+            for label in labels["input_ids"]
+        ]
+        
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+    
+    def compute_metrics(self, eval_pred):
+        """Compute metrics for evaluation"""
+        predictions, labels = eval_pred
+        
+        # Decode predictions
+        decoded_preds = self.tokenizer.batch_decode(
+            predictions, skip_special_tokens=True
+        )
+        
+        # Replace -100 in the labels as we can't decode them
+        labels = np.where(labels != -100, labels, self.tokenizer.pad_token_id)
+        decoded_labels = self.tokenizer.batch_decode(
+            labels, skip_special_tokens=True
+        )
+        
+        # Calculate exact match accuracy
+        exact_match = sum(
+            pred.strip().lower() == label.strip().lower()
+            for pred, label in zip(decoded_preds, decoded_labels)
+        ) / len(decoded_preds)
+        
+        logger.info(f"Exact match accuracy: {exact_match:.4f}")
+        
+        return {"exact_match": exact_match}
+    
+    def train(self, train_dataset, val_dataset, output_dir, hf_token=None, hub_username=None, num_epochs=5):
+        """Train the model"""
+        logger.info("Starting training preparation...")
+        
+        # Tokenize datasets
+        train_dataset = train_dataset.map(
+            self.preprocess_function,
+            batched=True,
+            remove_columns=['instruction', 'output', 'task_type', 'quality_score']
+        )
+        
+        val_dataset = val_dataset.map(
+            self.preprocess_function,
+            batched=True,
+            remove_columns=['instruction', 'output', 'task_type', 'quality_score']
+        )
+        
+        # Data collator
+        data_collator = DataCollatorForSeq2Seq(
+            self.tokenizer,
+            model=self.model,
+            label_pad_token_id=-100,
+            pad_to_multiple_of=8
+        )
+        
+        # Training arguments optimized for 4xL4 GPUs
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=num_epochs,
+            per_device_train_batch_size=32,  # L4 has 24GB, can handle larger batches
+            per_device_eval_batch_size=64,
+            gradient_accumulation_steps=1,
+            gradient_checkpointing=True,
+            fp16=True,  # Use mixed precision
+            optim="adamw_torch",
+            learning_rate=3e-4 if self.use_lora else 5e-5,
+            warmup_steps=500,
+            logging_steps=50,
+            evaluation_strategy="steps",
+            eval_steps=500,
+            save_strategy="steps",
+            save_steps=500,
+            load_best_model_at_end=True,
+            metric_for_best_model="exact_match",
+            greater_is_better=True,
+            push_to_hub=True if hf_token else False,
+            hub_model_id=f"{hub_username}/resume-normalizer-flan-t5" if hub_username else None,
+            hub_token=hf_token,
+            report_to=["tensorboard"],
+            ddp_find_unused_parameters=False,
+            dataloader_num_workers=4,
+            remove_unused_columns=False,
+        )
+        
+        # Create trainer
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=val_dataset,
+            tokenizer=self.tokenizer,
+            data_collator=data_collator,
+            compute_metrics=self.compute_metrics,
+            callbacks=[
+                EarlyStoppingCallback(early_stopping_patience=3)
+            ],
+        )
+        
+        logger.info("Starting training...")
+        
+        # Train
+        trainer.train()
+        
+        # Save model
+        logger.info("Saving model...")
+        if self.use_lora:
+            # Save LoRA adapter
+            self.model.save_pretrained(output_dir)
+            self.tokenizer.save_pretrained(output_dir)
+        else:
+            trainer.save_model(output_dir)
+            
+        # Push to hub if token provided
+        if hf_token and hub_username:
+            logger.info("Pushing model to HuggingFace Hub...")
+            trainer.push_to_hub(
+                commit_message="Final model trained on resume normalization data"
+            )
+            
+        logger.info("Training completed successfully!")
+        return trainer
+
+def main():
+    """Main training function to be called from app.py"""
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_path", type=str, required=True)
+    parser.add_argument("--model_size", type=str, default="base")
+    parser.add_argument("--hf_token", type=str, default=None)
+    parser.add_argument("--hub_username", type=str, default=None)
+    parser.add_argument("--num_epochs", type=int, default=5)
+    parser.add_argument("--use_lora", action="store_true")
+    args = parser.parse_args()
+    
+    # Set HF token if provided
+    if args.hf_token:
+        HfFolder.save_token(args.hf_token)
+    
+    # Select model based on size
+    model_name = "google/flan-t5-base" if args.model_size == "base" else "google/flan-t5-large"
+    
+    # Initialize trainer
+    trainer = ResumeNormalizationTrainer(
+        model_name=model_name,
+        use_lora=args.use_lora
+    )
+    
+    # Load dataset
+    train_dataset, val_dataset = trainer.load_dataset(args.data_path)
+    
+    # Train
+    output_dir = "./resume-normalizer-model"
+    trainer.train(
+        train_dataset=train_dataset,
+        val_dataset=val_dataset,
+        output_dir=output_dir,
+        hf_token=args.hf_token,
+        hub_username=args.hub_username,
+        num_epochs=args.num_epochs
+    )
+    
+    print("Training completed successfully!")
+    print(f"Model saved to: {output_dir}")
+    if args.hf_token and args.hub_username:
+        print(f"Model available at: https://huggingface.co/{args.hub_username}/resume-normalizer-flan-t5")
+
+if __name__ == "__main__":
+    main()