main.py - master - daniel-robotics/ML-Service-Deploy-Example

main.py
90 строк3 KB
Даниил Грабарь
create: .gitignore, Dockerfile, main.py, requirements.txt, .dockerignore
18 апр 2026, 09:43
18 апр 2026, 09:43130f8bc
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
import uvicorn
import logging


# Настройка логирования
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="Sentiment Analysis API",
    description="ML-сервис для анализа тональности текста на базе DistilBERT",
    version="1.0.0",
)

# Загружаем модель один раз при старте приложения
logger.info("Загрузка модели...")
classifier = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english",
)
logger.info("Модель загружена!")


# Схемы данных 
class PredictRequest(BaseModel):
    text: str

    model_config = {
        "json_schema_extra": {
            "examples": [{"text": "I love this product, it works great!"}]
        }
    }


class PredictResponse(BaseModel):
    text: str
    label: str      # POSITIVE / NEGATIVE
    score: float    # уверенность модели (0.0 — 1.0)


class BatchPredictRequest(BaseModel):
    texts: list[str]


class BatchPredictResponse(BaseModel):
    results: list[PredictResponse]


# Эндпоинты
@app.get("/health", summary="Проверка работоспособности сервиса")
def health():
    return {"status": "ok"}


@app.post("/predict", response_model=PredictResponse, summary="Предсказание для одного текста")
def predict(request: PredictRequest):
    if not request.text.strip():
        raise HTTPException(status_code=400, detail="Текст не может быть пустым")

    result = classifier(request.text)[0]
    return PredictResponse(
        text=request.text,
        label=result["label"],
        score=round(result["score"], 4),
    )


@app.post("/predict/batch", response_model=BatchPredictResponse, summary="Предсказание для нескольких текстов")
def predict_batch(request: BatchPredictRequest):
    if not request.texts:
        raise HTTPException(status_code=400, detail="Список текстов не может быть пустым")

    results = classifier(request.texts)
    return BatchPredictResponse(
        results=[
            PredictResponse(text=text, label=r["label"], score=round(r["score"], 4))
            for text, r in zip(request.texts, results)
        ]
    )


# http://localhost:8030/docs - документация Swagger UI
if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8030, reload=False)
    
# docker build -t sentiment-api .
# docker run -p 8030:8030 sentiment-api