sailokeshg/tech-stack-advisor
0
1import pandas as pd2from sklearn.preprocessing import LabelEncoder3from sklearn.tree import DecisionTreeClassifier4import pickle5 6# Sample synthetic data7data = {8 "project_type": ["Web App", "API", "ML App", "Real-time App", "Web App"],9 "team_size": [3, 2, 5, 6, 1],10 "perf_need": ["Medium", "Low", "Medium", "High", "Low"],11 "experience": ["Intermediate", "Beginner", "Expert", "Expert", "Beginner"],12 "stack": ["Django + PostgreSQL", "Flask + SQLite", "FastAPI + TensorFlow", "Node.js + Redis", "Django + SQLite"]13}14 15df = pd.DataFrame(data)16 17# Encode categorical variables18encoders = {}19for col in ["project_type", "perf_need", "experience", "stack"]:20 le = LabelEncoder()21 df[col] = le.fit_transform(df[col])22 encoders[col] = le23 24# Train model25X = df[["project_type", "team_size", "perf_need", "experience"]]26y = df["stack"]27model = DecisionTreeClassifier()28model.fit(X, y)29 30# Save model31with open("model.pkl", "wb") as f:32 pickle.dump(model, f)33 34# Save encoders for use in app35with open("encoders.pkl", "wb") as f:36 pickle.dump(encoders, f)37 