Update app.py

app.py

@@ -1,82 +1,97 @@
 import gradio as gr
-import joblib
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-import matplotlib.pyplot as plt
+import joblib
 import random
+import json
 
-# Load model
-model = joblib.load("model.joblib")  # Ensure your model was saved as a Pipeline
+# Load the model
+model = joblib.load("model.joblib")
 
-# Load base data for teams
-df = pd.read_csv("data.csv")  # Used to get team names
-teams = sorted(df['Team'].unique())
+# Team list (based on your dataset)
+teams = [f"team_{i}" for i in range(1, 13)]
 
-# Define possible randomized tasks
-def get_random_task():
-    tasks = [
-        {
-            'quarter': 'Q1', 'department': 'molding', 'day': 'Monday', 'no_of_workers': 40,
-            'incentive': 2.0, 'idle_time': 0.3, 'idle_men': 3, 'smv': 25.0, 'month': 3,
-            'day_of_week': 0, 'is_weekend': 0, 'log_wip': np.log1p(35), 'log_overtime': np.log1p(0.8),
-            'no_of_style_change': 0, 'targeted_productivity': 0.75
-        },
-        {
-            'quarter': 'Q3', 'department': 'packaging', 'day': 'Wednesday', 'no_of_workers': 52,
-            'incentive': 1.5, 'idle_time': 0.6, 'idle_men': 6, 'smv': 28.0, 'month': 8,
-            'day_of_week': 2, 'is_weekend': 0, 'log_wip': np.log1p(60), 'log_overtime': np.log1p(0.5),
-            'no_of_style_change': 1, 'targeted_productivity': 0.7
-        },
-        {
-            'quarter': 'Q4', 'department': 'quality', 'day': 'Friday', 'no_of_workers': 35,
-            'incentive': 3.0, 'idle_time': 0.1, 'idle_men': 1, 'smv': 32.0, 'month': 11,
-            'day_of_week': 4, 'is_weekend': 0, 'log_wip': np.log1p(20), 'log_overtime': np.log1p(1.2),
-            'no_of_style_change': 0, 'targeted_productivity': 0.85
-        }
-    ]
-    task = random.choice(tasks)
-    task['smv_per_worker'] = task['smv'] / task['no_of_workers']
-    task['effort_index'] = task['smv'] + task['incentive'] + 1.0 - task['idle_time']
-    return task
+# 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': 'Insecticide',
+        'TaskType': 'Labeling',
+        'OrderQuantity': 80,
+        'DeadlineDays': 2,
+        'ExperienceYears': 2,
+        'AvgTaskTime_Minutes': 45.0,
+        'ErrorRate': 0.12,
+        'TrainingHours': 10.0,
+        'DayNumber': 3
+    },
+    {
+        'ProductType': 'Soap',
+        'TaskType': 'Mixing',
+        'OrderQuantity': 200,
+        'DeadlineDays': 6,
+        'ExperienceYears': 4,
+        'AvgTaskTime_Minutes': 32.0,
+        'ErrorRate': 0.07,
+        'TrainingHours': 15.0,
+        'DayNumber': 4
+    },
+]
 
-# Function to predict productivity and return sorted chart and summary
+# Priority map
+priority_map = {'High': 3, 'Medium': 2, 'Low': 1}
+
+# Prediction function
 def predict():
-    task = get_random_task()
-    team_scores = []
+    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  # all samples are high priority
 
+    # Predict for each team
+    results = []
     for team in teams:
-        t = task.copy()
-        t['Team'] = team
-        task_df = pd.DataFrame([t])
-        pred = model.predict(task_df)[0]
-        team_scores.append((team, pred))
+        row = task.copy()
+        row['Team'] = team
+        input_df = pd.DataFrame([row])
+        pred = model.predict(input_df)[0]
+        results.append((team, round(pred, 4)))
 
-    df_scores = pd.DataFrame(team_scores, columns=["Team", "Predicted Productivity"])
-    df_scores = df_scores.sort_values(by="Predicted Productivity", ascending=False)
+    # Sort results
+    results.sort(key=lambda x: x[1], reverse=True)
 
     # Plot
-    fig, ax = plt.subplots(figsize=(8, 6))
-    ax.barh(df_scores['Team'], df_scores['Predicted Productivity'], color='skyblue')
+    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.set_title("Team Ranking for Random Task")
     ax.invert_yaxis()
     plt.tight_layout()
 
-    # Convert task to readable summary
-    task_summary = "\n".join([f"{k}: {v}" for k, v in task.items()])
-
-    return task_summary, fig
+    return fig, json.dumps(task, indent=2)
 
 # Gradio UI
 demo = gr.Interface(
     fn=predict,
     inputs=[],
-    outputs=[
-        gr.Textbox(label="📌 Random Task Details"),
-        gr.Plot(label="📈 Team Productivity Rankings")
-    ],
-    title="🧠 Team Assignment Predictor",
-    description="Click Submit to generate a random task and rank all teams based on predicted productivity."
+    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."
 )
 
-demo.launch()
+demo.launch()