CoolFace
Apppublic

sailokeshg/tech-stack-advisor

sourceHugging Faceapache-2.0updated 1y agoView on Hugging Face
0likes
train.py37 linesDownload Raw Back to root
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