Update label.py

models/label.py

@@ -32,19 +32,19 @@ class PredictionResponse(BaseModel):
             "example": {
                 "prediction": 1,
                 "confidence": 0.874,
-                "label": "MATCH",
+                "label": "apparie",
                 "probabilities": {
-                    "match": 0.874,
-                    "no_match": 0.126
+                    "non_apparie": 0.126,
+                    "apparie": 0.874
                 }
             }
         }
     )
 
-    prediction: int = Field(..., description="1 = match, 0 = no match")
+    prediction: int = Field(..., description="1 = apparie, 0 = non_apparie")
     confidence: float = Field(..., ge=0.0, le=1.0,
                               description="Confidence score of the prediction")
-    label: str = Field(..., description="MATCH or NO_MATCH")
+    label: str = Field(..., description="apparie or non_apparie")
     probabilities: Dict[str, float] = Field(
         ..., description="Dictionary of class probabilities"
     )
@@ -57,12 +57,24 @@ class BatchPredictionRequest(BaseModel):
             "example": {
                 "pairs": [
                     {
-                        "argument": "Schools should implement AI tools to support learning.",
-                        "key_point": "AI can improve student engagement."
+                        "argument": "School uniforms limit students' self-expression and creativity",
+                        "key_point": "Uniforms restrict personal freedom and individuality"
                     },
                     {
-                        "argument": "Governments must reduce plastic usage.",
-                        "key_point": "Plastic waste harms the environment."
+                        "argument": "Renewable energy creates more jobs than fossil fuel industries",
+                        "key_point": "Green energy generates employment opportunities"
+                    },
+                    {
+                        "argument": "We should invest more in renewable energy to combat climate change",
+                        "key_point": "Capital punishment violates human rights"
+                    },
+                    {
+                        "argument": "Online education provides flexibility and accessibility for all students",
+                        "key_point": "Digital learning offers convenient and inclusive education"
+                    },
+                    {
+                        "argument": "Vaccinations are essential for public health and disease prevention",
+                        "key_point": "Space exploration leads to scientific discoveries"
                     }
                 ]
             }
@@ -77,8 +89,72 @@ class BatchPredictionRequest(BaseModel):
 
 class BatchPredictionResponse(BaseModel):
     """Response model for batch key-point predictions"""
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "predictions": [
+                    {
+                        "prediction": 1,
+                        "confidence": 0.956,
+                        "label": "apparie",
+                        "probabilities": {
+                            "non_apparie": 0.044,
+                            "apparie": 0.956
+                        }
+                    },
+                    {
+                        "prediction": 1,
+                        "confidence": 0.892,
+                        "label": "apparie",
+                        "probabilities": {
+                            "non_apparie": 0.108,
+                            "apparie": 0.892
+                        }
+                    },
+                    {
+                        "prediction": 0,
+                        "confidence": 0.934,
+                        "label": "non_apparie",
+                        "probabilities": {
+                            "non_apparie": 0.934,
+                            "apparie": 0.066
+                        }
+                    },
+                    {
+                        "prediction": 1,
+                        "confidence": 0.878,
+                        "label": "apparie",
+                        "probabilities": {
+                            "non_apparie": 0.122,
+                            "apparie": 0.878
+                        }
+                    },
+                    {
+                        "prediction": 0,
+                        "confidence": 0.967,
+                        "label": "non_apparie",
+                        "probabilities": {
+                            "non_apparie": 0.967,
+                            "apparie": 0.033
+                        }
+                    }
+                ],
+                "total_processed": 5,
+                "summary": {
+                    "total_apparie": 3,
+                    "total_non_apparie": 2,
+                    "average_confidence": 0.9254
+                }
+            }
+        }
+    )
+
     predictions: List[PredictionResponse]
     total_processed: int = Field(..., description="Number of processed items")
+    summary: Dict[str, float] = Field(
+        ...,
+        description="Summary statistics of the batch prediction"
+    )
 
 
 class HealthResponse(BaseModel):
@@ -86,13 +162,49 @@ class HealthResponse(BaseModel):
     model_config = ConfigDict(
         json_schema_extra={
             "example": {
-                "status": "ok",
+                "status": "healthy",
                 "model_loaded": True,
-                "device": "cuda"
+                "device": "cpu",
+                "model_name": "NLP-Debater-Project/destlibert-keypoint-matching",
+                "timestamp": "2024-01-01T12:00:00Z"
             }
         }
     )
     
     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)")
+    device: str = Field(..., description="Device used for inference (cpu/cuda)")
+    model_name: Optional[str] = Field(None, description="Name of the loaded model")
+    timestamp: str = Field(..., description="Timestamp of the health check")
+
+
+class ModelInfoResponse(BaseModel):
+    """Detailed model information response"""
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "model_name": "NLP-Debater-Project/destlibert-keypoint-matching",
+                "device": "cpu",
+                "max_length": 256,
+                "num_labels": 2,
+                "loaded": True,
+                "performance": {
+                    "accuracy": 0.9285,
+                    "f1_score": 0.8836,
+                    "f1_apparie": 0.8113,
+                    "f1_non_apparie": 0.9559
+                },
+                "description": "DistilBERT model for key point - argument semantic matching"
+            }
+        }
+    )
+
+    model_name: str
+    device: str
+    max_length: int
+    num_labels: int
+    loaded: bool
+    performance: Dict[str, float] = Field(
+        ..., description="Model performance metrics"
+    )
+    description: str = Field(..., description="Model description")