File size: 2,038 Bytes
ac9ddbb |
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 |
"""API Schemas for Predict Request and Response."""
from enum import Enum
from pydantic import BaseModel, ConfigDict, ValidationError
class ProgrammingLanguage(str, Enum):
"""Programming languages supported for prediction."""
JAVA = "java"
PYTHON = "python"
PHARO = "pharo"
class ModelType(str, Enum):
"""Model types for prediction."""
SETFIT = "setfit"
RANDOM_FOREST = "random_forest"
TRANSFORMER = "transformer"
class PredictRequest(BaseModel):
"""Schema for Predict Request."""
text: str
language: ProgrammingLanguage
model_type: ModelType
model_config = ConfigDict(
json_schema_extra={
"example": {
"text": "This method calculates the average score.",
"language": "python",
"model_type": "transformer",
}
}
)
class PredictResponse(BaseModel):
"""Schema for Predict Response."""
label: str
score: float
""" Demonstration of object instantiation, printing,
and validation error handling with dummy use cases"""
if __name__ == "__main__":
print("\n--- 1. Object Instantiation & Printing ---")
valid_data = {
"text": "This method calculates the average score.",
"language": "java",
"model_type": "setfit",
}
# Instantiate the object
request = PredictRequest(**valid_data)
# Print object as dictionary (.model_dump() is Pydantic V2 syntax)
print(f"Valid Request Object: {request.model_dump()}")
print("\n--- 2. Handling Invalid Data ---")
try:
print("Attempting to create request with language='c++'...")
# This should fail because 'c++' is not in ProgrammingLanguage Enum
invalid_request = PredictRequest(
text="std::cout << 'Hello';", language="c++", model_type="setfit"
)
except ValidationError as e:
print("SUCCESS: Validation Error Caught!")
print(e.json())
|