models/label.py

"""Pydantic schemas for key-point matching prediction endpoints"""

from pydantic import BaseModel, Field, ConfigDict
from typing import List, Optional, Dict


class PredictionRequest(BaseModel):
    """Request model for single key-point/argument prediction"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "argument": "Climate change is accelerating due to industrial emissions.",
                "key_point": "Human industry contributes significantly to global warming."
            }
        }
    )

    argument: str = Field(
        ..., min_length=5, max_length=1000,
        description="The argument text to evaluate"
    )
    key_point: str = Field(
        ..., min_length=5, max_length=500,
        description="The key point used for comparison"
    )


class PredictionResponse(BaseModel):
    """Response model for single prediction"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "prediction": 1,
                "confidence": 0.874,
                "label": "MATCH",
                "probabilities": {
                    "match": 0.874,
                    "no_match": 0.126
                }
            }
        }
    )

    prediction: int = Field(..., description="1 = match, 0 = no match")
    confidence: float = Field(..., ge=0.0, le=1.0,
                              description="Confidence score of the prediction")
    label: str = Field(..., description="MATCH or NO_MATCH")
    probabilities: Dict[str, float] = Field(
        ..., description="Dictionary of class probabilities"
    )


class BatchPredictionRequest(BaseModel):
    """Request model for batch predictions"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "pairs": [
                    {
                        "argument": "Schools should implement AI tools to support learning.",
                        "key_point": "AI can improve student engagement."
                    },
                    {
                        "argument": "Governments must reduce plastic usage.",
                        "key_point": "Plastic waste harms the environment."
                    }
                ]
            }
        }
    )

    pairs: List[PredictionRequest] = Field(
        ..., max_length=100,
        description="List of argument-keypoint pairs (max 100)"
    )


class BatchPredictionResponse(BaseModel):
    """Response model for batch key-point predictions"""
    predictions: List[PredictionResponse]
    total_processed: int = Field(..., description="Number of processed items")


class HealthResponse(BaseModel):
    """Health check model for the API"""
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "status": "ok",
                "model_loaded": True,
                "device": "cuda"
            }
        }
    )
    
    status: str = Field(..., description="API health status")
    model_loaded: bool = Field(..., description="Whether the model is loaded")
    device: str = Field(..., description="Device used for inference (cpu/cuda)")