File size: 695 Bytes
5fc6e5d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from typing import Any, List

from pydantic import BaseModel, Field


# Input Schema
class PredictionRequest(BaseModel):
    texts: List[str] = Field(
        ...,
        description="List of code comments to classify",
        example=["public void main", "def init self"],
    )
    language: str = Field(
        ..., description="Programming language (java, python, pharo)", example="java"
    )


# Output Schema
class PredictionResponse(BaseModel):
    predictions: List[Any] = Field(..., description="List of predicted labels")
    labels: List[Any] = Field(..., description="List of human-readable labels")
    model_info: dict = Field(..., description="Metadata about the model used")