Calin Rada
init
f006f31 unverified
raw
history blame
2.9 kB
import logging.config
from enum import Enum
from typing import List, Union
from pydantic import BaseModel, validator
logger = logging.getLogger(__name__)
class PredictionResponse(BaseModel):
label: str
score: float
class BedType(str, Enum):
none = ""
single = "single"
double = "double"
queen = "queen"
king = "king"
california_king = "california_king"
bunk = "bunk"
sofa = "sofa"
rollaway = "rollaway"
futon = "futon"
class BedData(BaseModel):
type: Union[BedType, None] = None
count: Union[int, None] = None
@validator("count", pre=True, always=True)
def validate_count(cls, v):
return v if v and v > 0 else None
@validator("type", pre=True, always=True)
def validate_type(cls, v):
try:
return BedType(v.lower()) if v else None
except ValueError:
return None
@classmethod
def create_valid_bed(cls, type: Union[str, None], count: Union[int, None]):
if type and count and count > 0:
try:
type_enum = BedType(type.lower())
return cls(type=type_enum, count=count)
except ValueError:
logger.error(f"Invalid bed type: {type}")
return None
def __init__(self, **data):
variations = {
"individual": BedType.single,
"camaindividual": BedType.single,
"doble": BedType.double,
"camadoble": BedType.double,
"reina": BedType.queen,
"queenbed": BedType.queen,
"rey": BedType.king,
"californiaking": BedType.california_king,
"litera": BedType.bunk,
"sofá": BedType.sofa,
"sofacama": BedType.sofa,
"plegable": BedType.rollaway,
"futón": BedType.futon,
"twin": BedType.single,
"twinbed": BedType.single,
"singlebed": BedType.single,
"doublebed": BedType.double,
"largedouble": BedType.queen,
"extralargedouble": BedType.king,
"bunkbed": BedType.bunk,
"couch": BedType.sofa,
"airmattress": BedType.futon,
"floormattress": BedType.futon,
}
bed_type = data.get("type")
if bed_type:
normalized_bed_type = bed_type.replace(" ", "").lower()
data["type"] = variations.get(normalized_bed_type, bed_type)
super().__init__(**data)
class RoomData(BaseModel):
room_description: str
beds: Union[BedData, None]
class Predictions(BaseModel):
type: PredictionResponse
category: PredictionResponse
environment: List[PredictionResponse]
feature: List[PredictionResponse]
view: List[PredictionResponse]
language_detected: str
beds: List[BedData]
class AllPredictionsResponse(BaseModel):
predictions: Predictions