Spaces:
Sleeping
Sleeping
File size: 4,102 Bytes
725d1d2 6752520 725d1d2 6752520 725d1d2 6752520 725d1d2 6752520 725d1d2 6752520 | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | from typing import Literal, Optional
from pydantic import BaseModel, Field, StrictBool, model_validator, StrictFloat
from api.schema.input_helpers import (
JOB_CATEGORY_MODEL,
age_category_map,
education_map,
map_customer_job,
)
class Education(BaseModel):
type: Literal["school", "university", "course", "illiterate"]
level: Optional[Literal["primary", "middle", "high"]] = None
grade: Optional[int] = None
@model_validator(mode="before")
def validate_education(cls, values):
type_ = values.get("type")
level = values.get("level")
grade = values.get("grade")
if type_ == "school":
if level is None or grade is None:
raise ValueError(
"For type 'school', both level and grade must be provided"
)
else:
if level is not None or grade is not None:
raise ValueError(
f"For type '{type_}', level and grade must not be provided"
)
return values
class PersonalInfo(BaseModel):
age: int = Field(..., ge=0, le=120)
job: Literal[
"blue_collar",
"housemaid",
"services",
"admin",
"technician",
"management",
"self_employed",
"entrepreneur",
"unemployed",
"student",
]
marital: Literal["married", "single"]
education: Education
# Derived fields
job_category: Optional[Literal["cat1", "cat2", "cat3", "cat4", "other"]] = None
age_category: Optional[
Literal[
"struggling",
"stable",
"about to retire",
"old age",
"counting a last breathe",
]
] = None
education_str: Optional[str] = None
@model_validator(mode="before")
def derive_categories(cls, values):
job_input = values.get("job")
age_input = values.get("age")
edu_dict = values.get("education")
job_model_name = map_customer_job(job_input)
job_cat = JOB_CATEGORY_MODEL.get(job_model_name, "other")
age_cat = age_category_map(age_input)
edu_type = edu_dict.get("type") if edu_dict else None
edu_level = edu_dict.get("level") if edu_dict else None
edu_grade = edu_dict.get("grade") if edu_dict else None
values["education_str"] = education_map(
type=edu_type, level=edu_level, grade=edu_grade
)
values["job_category"] = job_cat
values["age_category"] = age_cat
return values
class FinancialInfo(BaseModel):
default: StrictBool
housing: StrictBool
loan: StrictBool
@model_validator(mode="after")
def transform_bool_to_int(cls, values):
values.default = int(values.default)
values.housing = int(values.housing)
values.loan = int(values.loan)
return values
class ContactInfo(BaseModel):
contact: Literal["cellular", "telephone"]
day_of_week: Literal["mon", "tue", "wed", "thu", "fri"]
month: Literal[
"mar",
"apr",
"may",
"jun",
"jul",
"aug",
"sep",
"oct",
"nov",
"dec",
]
class CampaignInfo(BaseModel):
campaign: int = Field(..., ge=0)
previous: int = Field(..., ge=0)
poutcome: Literal["success", "failure", "nonexistent"]
class MacroInfo(BaseModel):
emp_var_rate: StrictFloat = Field(..., alias="employment_variation_rate")
euribor3m: StrictFloat = Field(..., alias="euribor_3m_rate")
nr_employed: StrictFloat = Field(..., alias="number_employed")
cons_price_idx: StrictFloat = Field(..., alias="consumer_price_index")
cons_conf_idx: float = Field(..., alias="consumer_confidence_index")
class InputData(BaseModel):
personal_info: PersonalInfo
financial_info: FinancialInfo
contact_info: ContactInfo
campaign_info: CampaignInfo
macro_info: MacroInfo
|