main.py

# main.py
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, confusion_matrix
import matplotlib.pyplot as plt

# Load data
df = pd.read_csv('data/employee_data.csv')

# Encode categorical
le = LabelEncoder()
df['Department'] = le.fit_transform(df['Department'])
df['Performance'] = le.fit_transform(df['Performance'])

# Split
X = df.drop('Performance', axis=1)
y = df['Performance']

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Model
model = RandomForestClassifier()
model.fit(X_train, y_train)

# Prediction
y_pred = model.predict(X_test)

# Accuracy
acc = accuracy_score(y_test, y_pred)
print("Accuracy:", acc)

# Confusion Matrix
cm = confusion_matrix(y_test, y_pred)
print(cm)

# Plot
plt.imshow(cm)
plt.title("Confusion Matrix")
plt.colorbar()
plt.savefig("outputs/confusion_matrix.png")
plt.show()