Update app.py

app.py

@@ -1,221 +1,99 @@
-import streamlit as st
-import pandas as pd
+import gradio as gr
+import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import joblib
-import matplotlib.pyplot as plt
 import random
-
-# Set page config
-st.set_page_config(
-    page_title="Team Productivity Predictor",
-    page_icon="🏆",
-    layout="wide"
-)
+import json
 
 # Load the model
-@st.cache_resource
-def load_model():
-    try:
-        model_path = "model.joblib"
-        model = joblib.load(model_path)
-        return model
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        return None
+model = joblib.load("model.joblib")
+
+# Load the original dataset to extract real specialty mappings
+df = pd.read_csv("dataset.csv")  # <-- Make sure this file is uploaded to Hugging Face
+
+# Get actual teams and specialties
+teams = sorted(df['Team'].unique())
+specialty_map = dict(zip(df['Team'], df['Specialty']))
 
-def generate_sample_task():
-    """Generate random sample task data"""
-    
-    # Sample options for categorical features
-    product_types = ["Mothball", "Perfume", "Candle", "Soap"]
-    task_types = ["Packaging", "Molding", "Quality Check", "Boxing", "Shipping"]
-    
-    # Generate random task
-    product_type = random.choice(product_types)
-    task_type = random.choice(task_types)
-    order_quantity = random.randint(50, 500)
-    deadline_days = random.randint(2, 14)
-    experience_years = round(random.uniform(1.0, 15.0), 1)
-    avg_task_time = round(random.uniform(15.0, 60.0), 1)
-    error_rate = round(random.uniform(0.01, 0.15), 3)
-    training_hours = round(random.uniform(5.0, 50.0), 1)
-    day_number = random.randint(1, 7)
-    priority_level = random.randint(1, 5)
-    
-    # Calculate derived features
-    throughput_rate = order_quantity / avg_task_time
-    time_pressure = order_quantity / (deadline_days * avg_task_time)
-    
-    return {
-        'ProductType': product_type,
-        'TaskType': task_type,
-        'OrderQuantity': order_quantity,
-        'DeadlineDays': deadline_days,
-        'ExperienceYears': experience_years,
-        'AvgTaskTime_Minutes': avg_task_time,
-        'ErrorRate': error_rate,
-        'TrainingHours': training_hours,
-        'DayNumber': day_number,
-        'ThroughputRate': throughput_rate,
-        'TimePressure': time_pressure,
-        'PriorityLevel': priority_level
-    }
+# Sample task examples (all feature-engineered already)
+sample_tasks = [
+    {
+        'ProductType': 'Mothball',
+        'TaskType': 'Packaging',
+        'OrderQuantity': 120,
+        'DeadlineDays': 4,
+        'ExperienceYears': 6,
+        'AvgTaskTime_Minutes': 28.0,
+        'ErrorRate': 0.05,
+        'TrainingHours': 20.0,
+        'DayNumber': 2
+    },
+    {
+        'ProductType': 'Perfume',
+        'TaskType': 'Boxing',
+        'OrderQuantity': 80,
+        'DeadlineDays': 2,
+        'ExperienceYears': 2,
+        'AvgTaskTime_Minutes': 45.0,
+        'ErrorRate': 0.12,
+        'TrainingHours': 10.0,
+        'DayNumber': 3
+    },
+    {
+        'ProductType': 'Soap',
+        'TaskType': 'Molding',
+        'OrderQuantity': 200,
+        'DeadlineDays': 6,
+        'ExperienceYears': 4,
+        'AvgTaskTime_Minutes': 32.0,
+        'ErrorRate': 0.07,
+        'TrainingHours': 15.0,
+        'DayNumber': 4
+    },
+]
 
-def main():
-    st.title("🏆 Team Productivity Predictor")
-    st.write("Generate random task data and predict team productivity rankings")
-    
-    # Load model
-    model = load_model()
-    if model is None:
-        st.stop()
-    
-    # Team and Specialty mapping
-    teams_specialties = {
-        'Team_1': 'Molding',
-        'Team_2': 'Packaging',
-        'Team_3': 'Quality Check',
-        'Team_4': 'Molding',
-        'Team_5': 'Boxing',
-        'Team_6': 'Shipping',
-        'Team_7': 'Boxing',
-        'Team_8': 'Packaging'
-    }
-    
-    # Center the button
-    col1, col2, col3 = st.columns([1, 2, 1])
-    with col2:
-        predict_button = st.button("🎲 Generate Sample Task & Predict Rankings", 
-                                 type="primary", 
-                                 use_container_width=True)
-    
-    if predict_button:
-        # Generate sample task
-        sample_task = generate_sample_task()
-        
-        # Display task details
-        st.header("📋 Generated Task Details")
-        
-        # Create columns for task details
-        col1, col2, col3 = st.columns(3)
-        
-        with col1:
-            st.metric("🏷️ Product Type", sample_task['ProductType'])
-            st.metric("⚙️ Task Type", sample_task['TaskType'])
-            st.metric("📦 Order Quantity", f"{sample_task['OrderQuantity']:,}")
-            st.metric("⏰ Deadline Days", sample_task['DeadlineDays'])
-        
-        with col2:
-            st.metric("👨‍💼 Experience Years", f"{sample_task['ExperienceYears']}")
-            st.metric("⏱️ Avg Task Time (min)", f"{sample_task['AvgTaskTime_Minutes']}")
-            st.metric("❌ Error Rate", f"{sample_task['ErrorRate']:.3f}")
-            st.metric("📚 Training Hours", f"{sample_task['TrainingHours']}")
-        
-        with col3:
-            st.metric("📅 Day Number", sample_task['DayNumber'])
-            st.metric("⭐ Priority Level", sample_task['PriorityLevel'])
-            st.metric("🚀 Throughput Rate", f"{sample_task['ThroughputRate']:.2f}")
-            st.metric("⚡ Time Pressure", f"{sample_task['TimePressure']:.3f}")
-        
-        try:
-            # Build DataFrame with one row per team
-            rows = []
-            for team, specialty in teams_specialties.items():
-                r = sample_task.copy()
-                r['Team'] = team
-                r['Specialty'] = specialty
-                rows.append(r)
-            
-            test_df = pd.DataFrame(rows)
-            
-            # Make predictions
-            test_df['PredictedProductivity'] = model.predict(test_df)
-            ranked = test_df.sort_values('PredictedProductivity', ascending=False).reset_index(drop=True)
-            ranked['Rank'] = range(1, len(ranked) + 1)
-            
-            st.header("🏆 Team Productivity Rankings")
-            
-            # Display results in two columns
-            col1, col2 = st.columns([1, 1])
-            
-            with col1:
-                st.subheader("📊 Rankings Table")
-                
-                # Style the dataframe
-                display_df = ranked[['Rank', 'Team', 'Specialty', 'PredictedProductivity']].copy()
-                display_df['PredictedProductivity'] = display_df['PredictedProductivity'].round(4)
-                
-                # Add medal emojis for top 3
-                medals = {1: "🥇", 2: "🥈", 3: "🥉"}
-                display_df['Medal'] = display_df['Rank'].map(lambda x: medals.get(x, ""))
-                
-                # Reorder columns
-                display_df = display_df[['Medal', 'Rank', 'Team', 'Specialty', 'PredictedProductivity']]
-                
-                st.dataframe(
-                    display_df,
-                    use_container_width=True,
-                    hide_index=True
-                )
-                
-                # Show top 3 teams with special formatting
-                st.subheader("🎖️ Top 3 Teams")
-                for i in range(min(3, len(ranked))):
-                    team_data = ranked.iloc[i]
-                    medal = ["🥇", "🥈", "🥉"][i]
-                    st.success(f"{medal} **{team_data['Team']}** ({team_data['Specialty']}) - Score: {team_data['PredictedProductivity']:.4f}")
-            
-            with col2:
-                st.subheader("📈 Productivity Chart")
-                
-                # Create horizontal bar chart
-                fig, ax = plt.subplots(figsize=(8, 6))
-                
-                # Color gradient for bars
-                colors = plt.cm.viridis(np.linspace(0, 1, len(ranked)))
-                
-                bars = ax.barh(ranked['Team'], ranked['PredictedProductivity'], color=colors)
-                
-                # Customize chart
-                ax.invert_yaxis()
-                ax.set_xlabel('Predicted Productivity Score', fontsize=12)
-                ax.set_title('Team Productivity Rankings', fontsize=14, fontweight='bold')
-                ax.grid(axis='x', linestyle='--', alpha=0.3)
-                
-                # Add value labels on bars
-                for i, (bar, value) in enumerate(zip(bars, ranked['PredictedProductivity'])):
-                    ax.text(value + 0.001, bar.get_y() + bar.get_height()/2, 
-                           f'{value:.3f}', ha='left', va='center', fontsize=9)
-                
-                plt.tight_layout()
-                st.pyplot(fig)
-                plt.close()
-                
-        except Exception as e:
-            st.error(f"Error making predictions: {str(e)}")
-            st.write("Please check that your model file is compatible with the expected features.")
-    
-    else:
-        # Show instructions when no button is clicked
-        st.info("👆 Click the button above to generate a random task and see team productivity rankings!")
-        
-        # Show example of what teams are available
-        st.subheader("👥 Available Teams")
-        teams_df = pd.DataFrame([
-            {"Team": team, "Specialty": specialty} 
-            for team, specialty in {
-                'Team_1': 'Molding',
-                'Team_2': 'Packaging', 
-                'Team_3': 'Quality Check',
-                'Team_4': 'Molding',
-                'Team_5': 'Boxing',
-                'Team_6': 'Shipping',
-                'Team_7': 'Boxing',
-                'Team_8': 'Packaging'
-            }.items()
-        ])
-        
-        st.dataframe(teams_df, use_container_width=True, hide_index=True)
+# Prediction function
+def predict():
+    task = random.choice(sample_tasks)
+
+    # Feature engineering
+    task['ThroughputRate'] = task['OrderQuantity'] / task['AvgTaskTime_Minutes']
+    task['TimePressure'] = task['OrderQuantity'] / ((task['DeadlineDays'] or 1) * task['AvgTaskTime_Minutes'])
+    task['PriorityLevel'] = 3  # high priority
+
+    # Predict for each team
+    results = []
+    for team in teams:
+        row = task.copy()
+        row['Team'] = team
+        row['Specialty'] = specialty_map.get(team, "Unknown")
+        input_df = pd.DataFrame([row])
+        pred = model.predict(input_df)[0]
+        results.append((team, round(pred, 4)))
+
+    # Sort results
+    results.sort(key=lambda x: x[1], reverse=True)
+
+    # Plot
+    teams_sorted, scores_sorted = zip(*results)
+    fig, ax = plt.subplots(figsize=(10, 6))
+    ax.barh(teams_sorted, scores_sorted, color="skyblue")
+    ax.set_xlabel("Predicted Productivity")
+    ax.set_title("Team Ranking for Random Task")
+    ax.invert_yaxis()
+    plt.tight_layout()
+
+    return fig, json.dumps(task, indent=2)
+
+# Gradio UI
+demo = gr.Interface(
+    fn=predict,
+    inputs=[],
+    outputs=[gr.Plot(label="Team Productivity Rankings"), gr.Textbox(label="Task Details")],
+    live=False,
+    title="Team Productivity Predictor",
+    description="Click Generate to predict the best team for a randomly selected task."
+)
 
-if __name__ == "__main__":
-    main()
+demo.launch()