Added mypy rules and applied them.

Created logger and showcased basic return of info.

app/embeddings.py

@@ -1,5 +1,6 @@
 from transformers import AutoTokenizer, AutoModel
 from torch import Tensor
+from app.logger import logger
 
 model = AutoModel.from_pretrained("intfloat/multilingual-e5-large")
 tokenizer = AutoTokenizer.from_pretrained("intfloat/multilingual-e5-large")
@@ -28,6 +29,9 @@ def embed_text(texts: list[str]) -> list[list[float]]:
     batch_dict = tokenizer(
         texts, max_length=512, padding=True, truncation=True, return_tensors="pt"
     )
+    logger.info(
+        f"Tokenized {len(texts)} texts with number of tokens per text: {batch_dict['input_ids'].ne(tokenizer.pad_token_id).sum(dim=1).tolist()}"
+    )
     outputs = model(**batch_dict)
     embeddings = average_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
 

app/logger.py

@@ -1,5 +1,36 @@
 import logging
+from logging.config import dictConfig
 
-logging.basicConfig(
-    level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
-)
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "default": {
+            "format": "[%(asctime)s] [%(levelname)s] %(name)s: %(message)s",
+            "datefmt": "%Y-%m-%d %H:%M:%S",
+        },
+        "json": {
+            "format": (
+                '{"time": "%(asctime)s", '
+                '"level": "%(levelname)s", '
+                '"name": "%(name)s", '
+                '"message": "%(message)s"}'
+            ),
+            "datefmt": "%Y-%m-%d %H:%M:%S",
+        },
+    },
+    "handlers": {
+        "console": {
+            "class": "logging.StreamHandler",
+            "formatter": "default",
+        },
+    },
+    "root": {
+        "level": "INFO",
+        "handlers": ["console"],
+    },
+}
+
+dictConfig(LOGGING_CONFIG)
+
+logger = logging.getLogger("app")

app/main.py

@@ -1,21 +1,22 @@
 from fastapi import FastAPI, HTTPException
 from app.models import EmbedRequest, EmbedResponse
 from app.embeddings import embed_text
-import logging
+from app.logger import logger
 
 app = FastAPI(
     title="Embedding API",
     description="A simple API to generate text embeddings using Microsoft's `multilingual-e5-large` model.",
     version="1.0.0",
 )
-logger = logging.getLogger(__name__)
 
 
 @app.post("/embed", response_model=EmbedResponse)
-async def embed(request: EmbedRequest):
+async def embed(request: EmbedRequest) -> dict[str, list[list[float]]]:
     """Generate embeddings for a list of texts."""
+    logger.info("Generating embeddings...")
     try:
         vectors = embed_text(request.texts)
+        logger.info("Embeddings generated successfully!")
         return {"embeddings": vectors}
     except Exception as e:
         logger.exception("Error generating embeddings")
@@ -23,6 +24,6 @@ async def embed(request: EmbedRequest):
 
 
 @app.get("/health")
-async def health_check():
+async def health_check() -> dict[str, str]:
     """Health check endpoint."""
     return {"status": "ok"}

pyproject.toml

@@ -17,3 +17,27 @@ dependencies = [
     "transformers>=4.57.0",
     "uvicorn>=0.37.0",
 ]
+
+# https://quantlane.com/blog/type-checking-large-codebase/
+[tool.mypy]
+# Ensure full coverage
+disallow_untyped_calls = false
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_untyped_decorators = false
+check_untyped_defs = true
+
+# Restrict dynamic typing
+disallow_any_generics = false
+disallow_subclassing_any = false
+warn_return_any = false
+
+# Know exactly what you're doing
+warn_redundant_casts = true
+warn_unused_ignores = false
+warn_unused_configs = true
+warn_unreachable = true
+show_error_codes = true
+
+# Explicit is better than implicit
+no_implicit_optional = true

tests/test_api.py

@@ -4,7 +4,7 @@ from app.main import app
 client = TestClient(app)
 
 
-def test_embed():
+def test_embed() -> None:
     """Test the /embed endpoint with valid input."""
     response = client.post("/embed", json={"texts": ["query: Hello world"]})
     assert response.status_code == 200  # OK
@@ -13,13 +13,13 @@ def test_embed():
     assert len(data["embeddings"][0]) == 1024
 
 
-def test_embed_no_texts():
+def test_embed_no_texts() -> None:
     """Test the /embed endpoint with no texts provided."""
     response = client.post("/embed", json={})
     assert response.status_code == 422  # Unprocessable Entity
 
 
-def test_embed_long_text():
+def test_embed_long_text() -> None:
     """Test the /embed endpoint with a text longer than 2000 characters."""
     long_text = "query: " + "a" * 1994  # 2001 characters
     response = client.post("/embed", json={"texts": [long_text]})

tests/test_embeddings.py

@@ -3,7 +3,7 @@ import torch
 import pytest
 
 
-def test_average_pool_basic():
+def test_average_pool_basic() -> None:
     """Test average pooling produces correct shape and masking."""
     last_hidden_states = torch.tensor(
         [
@@ -29,7 +29,7 @@ def test_average_pool_basic():
     assert result.shape == (2, 2)
 
 
-def test_embed_text_valid():
+def test_embed_text_valid() -> None:
     """Test embedding returns correct number of vectors and dimensions."""
 
     texts = ["query: Hello world", "query: Hej verden"]
@@ -43,13 +43,13 @@ def test_embed_text_valid():
     assert len(embeddings[0]) == 1024
 
 
-def test_embed_text_empty_list():
+def test_embed_text_empty_list() -> None:
     """Should raise ValueError if no input texts."""
     with pytest.raises(ValueError, match="No input texts provided"):
         embed_text([])
 
 
-def test_embed_text_too_long():
+def test_embed_text_too_long() -> None:
     """Should raise ValueError for inputs exceeding 2000 characters."""
     too_long = ["query: " + "a" * 1994]  # 2001 characters
     with pytest.raises(ValueError, match="exceed the maximum length"):