File size: 1,077 Bytes
62af38c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
import pandas as pd
from sklearn.preprocessing import LabelEncoder
from sklearn.tree import DecisionTreeClassifier
import pickle
# Sample synthetic data
data = {
"project_type": ["Web App", "API", "ML App", "Real-time App", "Web App"],
"team_size": [3, 2, 5, 6, 1],
"perf_need": ["Medium", "Low", "Medium", "High", "Low"],
"experience": ["Intermediate", "Beginner", "Expert", "Expert", "Beginner"],
"stack": ["Django + PostgreSQL", "Flask + SQLite", "FastAPI + TensorFlow", "Node.js + Redis", "Django + SQLite"]
}
df = pd.DataFrame(data)
# Encode categorical variables
encoders = {}
for col in ["project_type", "perf_need", "experience", "stack"]:
le = LabelEncoder()
df[col] = le.fit_transform(df[col])
encoders[col] = le
# Train model
X = df[["project_type", "team_size", "perf_need", "experience"]]
y = df["stack"]
model = DecisionTreeClassifier()
model.fit(X, y)
# Save model
with open("model.pkl", "wb") as f:
pickle.dump(model, f)
# Save encoders for use in app
with open("encoders.pkl", "wb") as f:
pickle.dump(encoders, f)
|