move model classes to dedicated file

src/entity/model.py

@@ -0,0 +1,28 @@
+from typing import Literal, Optional, List
+from pydantic import BaseModel, Field
+
+
+class ModelInput(BaseModel):
+    p1_rank: int = Field(gt=0, default=1, description="The rank of the 1st player")
+    p2_rank: int = Field(gt=0, default=100, description="The rank of the 2nd player")
+    court: Literal['Outdoor', 'Indoor'] = Field(default='Outdoor', description="The type of court")
+    surface: Literal['Grass', 'Carpet', 'Clay', 'Hard'] = Field(default='Clay', description="The type of surface")
+    series: Literal['Grand Slam', 'Masters 1000', 'Masters', 'Masters Cup', 'ATP500', 'ATP250', 'International Gold', 'International'] = Field(default='Grand Slam', description="The series of the tournament")
+    p1_height: Optional[int] = Field(gt=0, default=180, description="The height of the 1st player in centimeters")
+    p2_height: Optional[int] = Field(gt=0, default=180, description="The height of the 2nd player in centimeters")
+    p1_weight: Optional[int] = Field(gt=0, default=80, description="The weight of the 1st player in kilograms")
+    p2_weight: Optional[int] = Field(gt=0, default=80, description="The weight of the 2nd player in kilograms")
+    p1_year_of_birth: Optional[int] = Field(gt=1950, default=1980, description="The year of birth of the 1st player")
+    p2_year_of_birth: Optional[int] = Field(gt=1950, default=1980, description="The year of birth of the 2nd player")
+    p1_play_hand: Literal['Right', 'Left'] = Field(default='Right', description="The play hand of the 1st player")
+    p2_play_hand: Literal['Right', 'Left'] = Field(default='Right', description="The play hand of the 2nd player")
+    p1_back_hand: int = Field(default=1, ge=1, le=2, description="The back hand of the 1st player. 1 for one-handed, 2 for two-handed")
+    p2_back_hand: int = Field(default=1, ge=1, le=2, description="The back hand of the 2nd player. 1 for one-handed, 2 for two-handed")
+    p1_pro_year: Optional[int] = Field(gt=1970, default=2000, description="The year the 1st player turned pro")
+    p2_pro_year: Optional[int] = Field(gt=1970, default=2000, description="The year the 2nd player turned pro")
+    model: Optional[str] = Field(default='LogisticRegression', description="The name of the model to use for prediction")
+    version: Optional[str] = Field(default='latest', description="The version of the model to use for prediction")
+
+class ModelOutput(BaseModel):
+    result: int = Field(description="The prediction result. 1 if player 1 is expected to win, 0 otherwise.", json_schema_extra={"example": "1"})
+    prob: List[float] = Field(description="Probability of [defeat, victory] of player 1.", json_schema_extra={"example": "[0.15, 0.85]"})

src/main.py

@@ -2,7 +2,7 @@ import os
 import joblib
 import logging
 import secrets
-from typing import Generator, Literal, Optional, Annotated
+from typing import Generator, Optional, Annotated
 from fastapi import (
     FastAPI,
     Request,
@@ -19,10 +19,11 @@ from starlette.status import (
     HTTP_403_FORBIDDEN,
     HTTP_404_NOT_FOUND,
     HTTP_503_SERVICE_UNAVAILABLE)
-from pydantic import BaseModel, Field
+from pydantic import Field
 from dotenv import load_dotenv
 from mlflow.exceptions import RestException
 
+from src.entity.model import ModelInput, ModelOutput
 from src.service.model import (
     run_experiment,
     predict,
@@ -93,31 +94,6 @@ async def run_xp(background_tasks: BackgroundTasks,
     
     return {"message": "Experiment scheduled"}
 
-class ModelInput(BaseModel):
-    p1_rank: int = Field(gt=0, default=1, description="The rank of the 1st player")
-    p2_rank: int = Field(gt=0, default=100, description="The rank of the 2nd player")
-    court: Literal['Outdoor', 'Indoor'] = Field(default='Outdoor', description="The type of court")
-    surface: Literal['Grass', 'Carpet', 'Clay', 'Hard'] = Field(default='Clay', description="The type of surface")
-    series: Literal['Grand Slam', 'Masters 1000', 'Masters', 'Masters Cup', 'ATP500', 'ATP250', 'International Gold', 'International'] = Field(default='Grand Slam', description="The series of the tournament")
-    p1_height: Optional[int] = Field(gt=0, default=180, description="The height of the 1st player in centimeters")
-    p2_height: Optional[int] = Field(gt=0, default=180, description="The height of the 2nd player in centimeters")
-    p1_weight: Optional[int] = Field(gt=0, default=80, description="The weight of the 1st player in kilograms")
-    p2_weight: Optional[int] = Field(gt=0, default=80, description="The weight of the 2nd player in kilograms")
-    p1_year_of_birth: Optional[int] = Field(gt=1950, default=1980, description="The year of birth of the 1st player")
-    p2_year_of_birth: Optional[int] = Field(gt=1950, default=1980, description="The year of birth of the 2nd player")
-    p1_play_hand: Literal['Right', 'Left'] = Field(default='Right', description="The play hand of the 1st player")
-    p2_play_hand: Literal['Right', 'Left'] = Field(default='Right', description="The play hand of the 2nd player")
-    p1_back_hand: int = Field(default=1, ge=1, le=2, description="The back hand of the 1st player. 1 for one-handed, 2 for two-handed")
-    p2_back_hand: int = Field(default=1, ge=1, le=2, description="The back hand of the 2nd player. 1 for one-handed, 2 for two-handed")
-    p1_pro_year: Optional[int] = Field(gt=1970, default=2000, description="The year the 1st player turned pro")
-    p2_pro_year: Optional[int] = Field(gt=1970, default=2000, description="The year the 2nd player turned pro")
-    model: Optional[str] = Field(default='LogisticRegression', description="The name of the model to use for prediction")
-    version: Optional[str] = Field(default='latest', description="The version of the model to use for prediction")
-
-class ModelOutput(BaseModel):
-    result: int = Field(description="The prediction result. 1 if player 1 is expected to win, 0 otherwise.", json_schema_extra={"example": "1"})
-    prob: list[float] = Field(description="Probability of [defeat, victory] of player 1.", json_schema_extra={"example": "[0.15, 0.85]"})
-
 @app.get("/predict",
          tags=["model"],
          description="Predict the outcome of a tennis match",