File size: 862 Bytes
c01955c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
from pydantic import BaseModel

class Train(BaseModel):
    model_name: str
    model_params: dict
    type: str
    make_dataset: dict

    class Config:
        json_schema_extra = {
            "examples": [
                {
                    "model_name": "LinearRegression",
                    "model_params": {"fit_intercept": True},
                    "type": "regression",
                    "make_dataset": {"n_samples": 100, "n_features": 2, "noise": 0.1}
                },
                {
                    "model_name": "RandomForestClassifier",
                    "model_params": {"n_estimators": 100, "max_depth": 5},
                    "type": "classification",
                    "make_dataset": {"n_samples": 100, "n_features": 4, "n_informative": 2, "n_classes": 2}
                },
                
            ]
        }