Initial commit

train/schemas.py

@@ -0,0 +1,37 @@
+from pydantic import BaseModel, Field
+
+
+class TrainRequest(BaseModel):
+    epochs:     int   = Field(10,    description="학습 에폭 수")
+    lr:         float = Field(1e-5,  description="학습률")
+    batch_size: int   = Field(4,     description="배치 사이즈")
+    ewc_lambda: float = Field(1000., description="EWC 패널티 강도 (클수록 기존 지식 보호)")
+    strategy:   str   = Field("full",description="auto | full | ewc | replay")
+
+
+class TrainStatus(BaseModel):
+    running:  bool
+    epoch:    int
+    total:    int
+    loss:     float
+    val_acc:  float = 0.0
+    message:  str
+    batch:         int = 0    # 현재 배치
+    total_batches: int = 0    # 전체 배치 수
+
+class PredictResult(BaseModel):
+    filename:       str
+    label:          str | None = None
+    confidence:     float | None = None
+    all_probs:      dict | None = None
+    key_tokens:     list[str] | None = None
+    ocr_word_count: int | None = None
+    total_pages:    int | None = None
+    per_page:       list | None = None
+    error:          str | None = None
+
+
+class DatasetStatus(BaseModel):
+    total_samples: int
+    by_label:      dict
+    label2id:      dict