labor-arbitration-system/nlp-service/app/main.py

"""
NLP microservice for labor arbitration element extraction.
Serves the fine-tuned Chinese RoBERTa multi-task model.
"""

from fastapi import FastAPI

from app.schemas import ExtractRequest
from app.services.extractor import ChineseLegalElementExtractor

app = FastAPI(title="Labor NLP Service", version="0.2.0")

# Initialize extractor (loads model on first use)
_extractor: ChineseLegalElementExtractor | None = None
_model_available: bool | None = None


def get_extractor() -> ChineseLegalElementExtractor:
    global _extractor
    if _extractor is None:
        _extractor = ChineseLegalElementExtractor()
    return _extractor


@app.get("/health")
def health():
    global _model_available
    ext = get_extractor()
    if _model_available is None:
        _model_available = ext._use_model
    return {
        "status": "ok",
        "model_available": _model_available,
        "mode": "bert_multi_task" if _model_available else "rule_fallback",
    }


@app.post("/extract")
def extract_elements(payload: ExtractRequest):
    ext = get_extractor()
    result = ext.extract(payload.text)
    return result


@app.post("/extract/compare")
def extract_compare(payload: ExtractRequest):
    """Return extraction results with both model and rule-based methods."""
    ext = get_extractor()
    model_result = None
    rule_result = None

    # Try model extraction
    if ext._use_model:
        try:
            model_result = ext._extract_with_model(payload.text)
        except Exception:
            pass

    # Rule-based extraction
    rule_result = ext._extract_with_rules(payload.text)

    return {
        "model_result": model_result,
        "rule_result": rule_result,
        "model_available": ext._use_model,
    }