Update app.py

app.py

@@ -1,85 +1,82 @@
 import gradio as gr
-import joblib
 import pandas as pd
 import numpy as np
+import matplotlib.pyplot as plt
+import joblib
 
-# 1) Load your trained pipeline
+# Load model
 model = joblib.load("model.joblib")
 
-# 2) Load team→specialty mapping
-df = pd.read_csv("dataset.csv")
-team_specialty = df[['Team','Specialty']].drop_duplicates().set_index('Team')['Specialty'].to_dict()
-teams = list(team_specialty.keys())
-
-# 3) Define several example task scenarios
-SCENARIOS = {
-    "Mothball → Packaging": {
-        'ProductType':         'Mothball',
-        'TaskType':            'Packaging',
-        'OrderQuantity':       100,
-        'DeadlineDays':        3,
-        'ExperienceYears':     5,
-        'AvgTaskTime_Minutes': 25.4,
-        'ErrorRate':           0.05,
-        'TrainingHours':       20.0,
-        'DayNumber':           2,
-        'ThroughputRate':      100 / 25.4,
-        'TimePressure':        100 / (3 * 25.4),
-        'PriorityLevel':       3
-    },
+# Define example tasks
+example_tasks = {
     "Soap → Molding": {
-        'ProductType':         'Soap',
-        'TaskType':            'Molding',
-        'OrderQuantity':       250,
-        'DeadlineDays':        5,
-        'ExperienceYears':     4,
-        'AvgTaskTime_Minutes': 40.0,
-        'ErrorRate':           0.10,
-        'TrainingHours':       15.0,
-        'DayNumber':           1,
-        'ThroughputRate':      250 / 40.0,
-        'TimePressure':        250 / (5 * 40.0),
-        'PriorityLevel':       2
+        'quarter': 'Q2', 'department': 'sewing', 'day': 'Monday',
+        'no_of_workers': 48, 'incentive': 2.5, 'idle_time': 0.3,
+        'idle_men': 4, 'smv': 30.0, 'month': 5, 'day_of_week': 0,
+        'is_weekend': 0, 'smv_per_worker': 30.0/48,
+        'effort_index': 30.0 + 2.5 + 1.0 - 0.3,
+        'log_wip': np.log1p(50), 'log_overtime': np.log1p(1.0),
+        'no_of_style_change': 0, 'targeted_productivity': 0.75
+    },
+    "Mothball → Packaging": {
+        'quarter': 'Q3', 'department': 'finishing', 'day': 'Wednesday',
+        'no_of_workers': 60, 'incentive': 3.0, 'idle_time': 0.5,
+        'idle_men': 6, 'smv': 25.0, 'month': 7, 'day_of_week': 2,
+        'is_weekend': 0, 'smv_per_worker': 25.0/60,
+        'effort_index': 25.0 + 3.0 + 1.0 - 0.5,
+        'log_wip': np.log1p(40), 'log_overtime': np.log1p(2.0),
+        'no_of_style_change': 1, 'targeted_productivity': 0.68
     },
     "Candle → QualityCheck": {
-        'ProductType':         'Candle',
-        'TaskType':            'QualityCheck',
-        'OrderQuantity':       150,
-        'DeadlineDays':        2,
-        'ExperienceYears':     6,
-        'AvgTaskTime_Minutes': 20.0,
-        'ErrorRate':           0.02,
-        'TrainingHours':       30.0,
-        'DayNumber':           3,
-        'ThroughputRate':      150 / 20.0,
-        'TimePressure':        150 / (2 * 20.0),
-        'PriorityLevel':       3
+        'quarter': 'Q1', 'department': 'quality', 'day': 'Friday',
+        'no_of_workers': 42, 'incentive': 1.5, 'idle_time': 0.1,
+        'idle_men': 2, 'smv': 20.0, 'month': 3, 'day_of_week': 4,
+        'is_weekend': 0, 'smv_per_worker': 20.0/42,
+        'effort_index': 20.0 + 1.5 + 1.0 - 0.1,
+        'log_wip': np.log1p(70), 'log_overtime': np.log1p(1.5),
+        'no_of_style_change': 0, 'targeted_productivity': 0.80
     }
 }
 
-def generate_rankings(scenario_name):
-    base = SCENARIOS[scenario_name]
-    rows = []
+# Define team list
+teams = [f"team_{i}" for i in range(1, 13)]
+
+# Prediction function
+def predict_productivity(task_name):
+    base = example_tasks[task_name]
+    predictions = []
+
     for team in teams:
-        r = base.copy()
-        r['Team'] = team
-        r['Specialty'] = team_specialty[team]
-        rows.append(r)
-    df_tasks = pd.DataFrame(rows)
-    df_tasks['PredictedProd'] = np.round(model.predict(df_tasks), 2)
-    return df_tasks[['Team', 'PredictedProd']].sort_values(
-        'PredictedProd', ascending=False
-    ).reset_index(drop=True)
+        row = base.copy()
+        row['Team'] = team
+        df_row = pd.DataFrame([row])
+        pred = model.predict(df_row)[0]
+        predictions.append((team, round(pred, 5)))
+
+    result_df = pd.DataFrame(predictions, columns=["Team", "Predicted Productivity"])
+    result_df = result_df.sort_values(by="Predicted Productivity", ascending=False)
+
+    # Plot bar chart
+    fig, ax = plt.subplots(figsize=(10, 6))
+    ax.barh(result_df["Team"], result_df["Predicted Productivity"], color="skyblue")
+    ax.set_xlabel("Predicted Productivity")
+    ax.set_title(f"Team Ranking for: {task_name}")
+    ax.invert_yaxis()
+    plt.tight_layout()
+
+    return result_df, fig
 
-with gr.Blocks() as demo:
-    gr.Markdown("## 🚀 Team Productivity Rankings")
-    scenario = gr.Dropdown(
-        choices=list(SCENARIOS.keys()),
-        label="Select a task scenario",
-        value=list(SCENARIOS.keys())[0]
-    )
-    btn = gr.Button("Generate Rankings")
-    table = gr.Dataframe(headers=["Team", "PredictedProd"])
-    btn.click(fn=generate_rankings, inputs=scenario, outputs=table)
+# Gradio UI
+demo = gr.Interface(
+    fn=predict_productivity,
+    inputs=gr.Dropdown(choices=list(example_tasks.keys()), label="Select Task Scenario"),
+    outputs=[
+        gr.Dataframe(label="🏆 Team Productivity Rankings:"),
+        gr.Plot(label="📈 Team Productivity Graph")
+    ],
+    title="👷 Team Assignment Predictor",
+    description="Select a task to predict productivity scores for all teams. The best team will be at the top!"
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()