api/schemas.py

"""
Pydantic Schemas for Bilingual Summarization REST API
"""

from typing import List, Optional, Dict, Any, Union
from pydantic import BaseModel, Field


# --- Health & System Info ---
class HealthResponse(BaseModel):
    status: str = Field(..., example="healthy")
    version: str = Field(..., example="1.0.0")
    device: str = Field(..., example="cpu")
    cuda_available: bool = Field(..., example=False)
    loaded_checkpoints: Dict[str, bool] = Field(
        ...,
        example={"seq2seq_arabic": True, "seq2seq_english": True, "rich_seq2seq_ar": True}
    )


class ModelDetails(BaseModel):
    id: str
    name: str
    paradigm: str  # "extractive" or "abstractive"
    supported_languages: List[str]
    description: str
    checkpoint_available: Optional[bool] = None


class ModelsListResponse(BaseModel):
    total_models: int
    models: List[ModelDetails]


# --- Language Detection ---
class LanguageDetectionRequest(BaseModel):
    text: str = Field(..., description="Raw text string to detect language for", min_length=1)

    model_config = {
        "json_schema_extra": {
            "example": {
                "text": "اعلنت وكالة ناسا الفضائية عن خطط طموحة للعودة الى سطح القمر في اطار برنامج ارتيميس."
            }
        }
    }


class LanguageDetectionResponse(BaseModel):
    detected_language: str = Field(..., example="ar")
    language_name: str = Field(..., example="Arabic")
    confidence: float = Field(..., example=0.98)
    script_breakdown: Dict[str, float] = Field(
        ...,
        example={"arabic": 0.95, "latin": 0.05, "other": 0.0}
    )


# --- Tokenization & Preprocessing ---
class TokenizeRequest(BaseModel):
    text: str = Field(..., min_length=1)
    lang: str = Field("auto", description="Language: 'auto', 'ar', or 'en'")
    remove_stopwords: bool = Field(False, description="Filter out stop words")
    stem: bool = Field(False, description="Apply word stemming")

    model_config = {
        "json_schema_extra": {
            "example": {
                "text": "الذكاء الاصطناعي يُحدث ثورة في معالجة اللغات الطبيعية. هل يمكن تلخيص النصوص بدقة؟",
                "lang": "ar",
                "remove_stopwords": False,
                "stem": False
            }
        }
    }


class TokenizeResponse(BaseModel):
    detected_language: str = Field(..., example="ar")
    sentence_count: int = Field(..., example=2)
    word_count: int = Field(..., example=13)
    sentences: List[str]
    tokens: List[str]
    normalized_text: str


# --- Summarization ---
class SummarizeRequest(BaseModel):
    text: str = Field(..., min_length=1, description="Source document text to summarize")
    mode: str = Field("extractive", description="Paradigm: 'extractive', 'abstractive', or 'hybrid'")
    method: str = Field(
        "textrank",
        description="Algorithm: 'textrank', 'lsa', 'hybrid', 'seq2seq', 'transformer'"
    )
    lang: str = Field("auto", description="'auto', 'ar', or 'en'")
    sentences: int = Field(3, ge=1, le=50, description="Target number of sentences (for extractive)")
    ratio: Optional[float] = Field(None, ge=0.05, le=0.95, description="Compression ratio (e.g. 0.3 for 30%)")
    beam_width: int = Field(3, ge=1, le=8, description="Beam search width for Seq2Seq decoding")
    reference_summary: Optional[str] = Field(
        None,
        description="Optional gold-standard reference summary to automatically compute ROUGE & BLEU"
    )

    model_config = {
        "json_schema_extra": {
            "example": {
                "text": "اعلنت وكالة ناسا الفضائية الامريكية عن خطط طموحة للعودة الى سطح القمر في اطار برنامج ارتيميس حيث ستشمل البعثة اول امراة واول شخص من ذوي البشرة الداكنة يمشيان على سطح القمر ومن المقرر ان تنطلق البعثة قبل نهاية عام الفين وخمسة وعشرين.",
                "mode": "abstractive",
                "method": "seq2seq",
                "lang": "ar",
                "beam_width": 3,
                "reference_summary": "ناسا تعلن خطط العودة للقمر عبر برنامج ارتيميس قبل نهاية 2025"
            }
        }
    }


class SummaryMetrics(BaseModel):
    original_words: int
    summary_words: int
    compression_ratio: float
    reduction_percentage: float
    type_token_ratio: float
    rouge_1_f1: Optional[float] = None
    rouge_2_f1: Optional[float] = None
    rouge_l_f1: Optional[float] = None
    bleu_1: Optional[float] = None
    bleu_2: Optional[float] = None
    bleu_cumulative: Optional[float] = None


class SentenceScoreItem(BaseModel):
    sentence: str
    score: float

class SummarizeResponse(BaseModel):
    status: str = "success"
    detected_language: str
    mode: str
    method: str
    summary: str
    latency_ms: float
    metrics: SummaryMetrics
    selected_indices: Optional[List[int]] = None
    sentence_scores: Optional[List[SentenceScoreItem]] = None
    note: Optional[str] = None


# --- Evaluation Request & Response ---
class EvaluationRequest(BaseModel):
    original_text: str = Field(..., min_length=1)
    generated_summary: str = Field(..., min_length=1)
    reference_summary: Optional[str] = Field(None, description="Ground truth reference summary")
    lang: str = Field("auto", description="'auto', 'ar', or 'en'")

    model_config = {
        "json_schema_extra": {
            "example": {
                "original_text": "Global semiconductor shortages are easing according to the latest industry report with chip inventories returning to normal.",
                "generated_summary": "Semiconductor shortage eases as chip inventories return to normal levels.",
                "reference_summary": "Global chip shortage easing with inventories normalizing across industry.",
                "lang": "en"
            }
        }
    }


class EvaluationResponse(BaseModel):
    language: str
    compression_ratio: float
    reduction_percentage: float
    original_words: int
    summary_words: int
    type_token_ratio: float
    rouge_1: Optional[Dict[str, float]] = None  # {precision, recall, f1}
    rouge_2: Optional[Dict[str, float]] = None
    rouge_l: Optional[Dict[str, float]] = None
    bleu_1: Optional[float] = None
    bleu_2: Optional[float] = None
    bleu_3: Optional[float] = None
    bleu_4: Optional[float] = None
    bleu_cumulative: Optional[float] = None