feat: implement init api sample

.gitignore

@@ -164,3 +164,4 @@ cython_debug/
 # Project
 healthcare-dataset-stroke-data.csv
 catboost_info/
+cvepred_model.pkl

Dockerfile

@@ -23,8 +23,14 @@ RUN poetry install --only main
 
 COPY cvepred /app/cvepred
 
+FROM base AS train
+
+RUN poetry run python3 -m cvepred.model
+
 FROM base AS api
 
 RUN poetry install --only main,api
 
+COPY --from=train /app/cvepred/cvepred_model.pkl /app/cvepred/cvepred_model.pkl
+
 CMD ["poetry", "run", "fastapi", "run", "cvepred/api.py"]

cvepred/api.py

@@ -1,8 +1,37 @@
 from fastapi import FastAPI
+from fastapi.responses import RedirectResponse
+from cvepred.dataset import create_dataset, split_datasets, RANDOM_STATE
+from cvepred.api_models import CveModels
+from pycaret.classification import setup
 
 app = FastAPI()
 
+samples = create_dataset()
+(df_train, df_test) = split_datasets(samples)
+
+s = setup(
+    data=df_train,
+    test_data=df_test,
+    target="hasExploit",
+    session_id=RANDOM_STATE,
+    preprocess=False,
+)
+model = s.load_model("cvepred_model")
+
+
 @app.get("/")
-def read_root():
-    return {"Hello": "World"}
+def redirect_to_docs():
+    return RedirectResponse(url="/docs")
+
+
+@app.get("/sample")
+def sample() -> CveModels:
+    data = samples.sample(n=20).drop(columns=["hasExploit"]).to_dict("records")
+
+    return CveModels.model_validate({"data": data})
+
 
+@app.post("/predict")
+def predict(inputs: CveModels):
+    # s.predict_model(model, )
+    raise Exception("Not implemented yet")

cvepred/api_models.py

@@ -0,0 +1,28 @@
+from pydantic import BaseModel
+import typing as ty
+
+# attackVector             category
+# attackComplexity         category
+# privilegesRequired       category
+# userInteraction          category
+# scope                    category
+# confidentialityImpact    category
+# integrityImpact          category
+# availabilityImpact       category
+# hasExploit                   bool
+
+
+class CveModel(BaseModel):
+    attackVector: str
+    attackComplexity: str
+    privilegesRequired: str
+    userInteraction: str
+    scope: str
+    confidentialityImpact: str
+    integrityImpact: str
+    availabilityImpact: str
+    # hasExploit: bool
+
+
+class CveModels(BaseModel):
+    data: ty.List[CveModel]

cvepred/dataset.py

@@ -113,3 +113,6 @@ def create_train_datasets():
     print("Datasets created")
 
     return df_balanced_train, xdf_balanced_test
+
+if __name__ == '__main__':
+    print(create_dataset().dtypes)

cvepred/model.py

@@ -16,4 +16,5 @@
 cvepred_model = s.create_model("rf", verbose=False)
 
 if __name__ == "__main__":
+    s.save_model(cvepred_model, "cvepred_model")
     print(s.pull())

pyproject.toml

@@ -13,7 +13,6 @@ kagglehub = "^0.3.3"
 pycaret = "^3.3.2"
 kaleido = "0.2.1"
 matplotlib = "<3.8.0"
-catboost = "^1.2.7"
 imbalanced-learn = "^0.12.4"
 httpx = "^0.27.2"