Use a prediciton dataclass

This commit is contained in:
Jacob Windsor 2025-02-19 13:27:02 +01:00
parent d6d2325b07
commit 6d0b501620
2 changed files with 11 additions and 14 deletions

View File

@ -1,12 +0,0 @@
from pydantic import BaseModel, Field
from typing import List
class HousePredictionInput(BaseModel):
square_feet: float = Field(..., gt=0, description="Square footage of the house")
bedrooms: int = Field(..., ge=1, description="Number of bedrooms")
bathrooms: float = Field(..., gt=0, description="Number of bathrooms")
class HousePredictionOutput(BaseModel):
predicted_price: float
confidence_score: float
similar_listings: List[float]

View File

@ -1,5 +1,14 @@
import random
from typing import List, Tuple
from dataclasses import dataclass
@dataclass
class Prediction:
predicted_price: float
confidence_score: float
similar_listings: List[float]
class HousePricePredictor:
"""
@ -11,7 +20,7 @@ class HousePricePredictor:
# Mock initialization - in reality would load model weights
pass
def predict(self, square_feet: float, bedrooms: int, bathrooms: float) -> Tuple[float, float, List[float]]:
def predict(self, square_feet: float, bedrooms: int, bathrooms: float) -> Prediction:
"""
Mock prediction method that returns:
- predicted price
@ -32,4 +41,4 @@ class HousePricePredictor:
for _ in range(3)
]
return predicted_price, confidence_score, similar_listings
return Prediction(predicted_price, confidence_score, similar_listings)