feat: update translation API to support generation parameters

- Modified request and response formats to include generation parameters for translation.
- Updated documentation to reflect changes in API usage and added details about generation parameters.

Author: ggwozdz90

docs/DOCKERHUB.md

@@ -74,20 +74,32 @@ Available on [Docker Hub](https://hub.docker.com/r/ggwozdz/translation-api):
 - Request:
 
     ```bash
-    curl -X POST "http://localhost:8000/translate" \
-        -F "text=Hello, how are you?" \
-        -F "source_language=en_US"
-        -F "target_language=pl_PL"
+    curl -X 'POST' \
+      'http://127.0.0.1:8000/translate' \
+      -H 'accept: application/json' \
+      -H 'Content-Type: application/json' \
+      -d '{
+      "text_to_translate": "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.",
+      "source_language": "en_US",
+      "target_language": "pl_PL",
+      "generation_parameters": { "max_length": 10240, "num_beams": 10 }
+    }'
     ```
 
 - Response:
 
     ```json
     {
-      "content": "Cześć, jak się masz?",
+      "translation": "Wieża Eiffla ma wysokość 324 metrów, mniej więcej taką samą wysokość jak 81-piętrowy budynek, i jest najwyższą budowlą w Paryżu. Jego podstawa jest kwadratowa, mierząc 125 metrów na każdej stronie. Podczas jej budowy Wieża Eiffla przekroczyła Pomnik Waszyngtonu, stając się najwyższą budowlą stworzoną przez człowieka na świecie, tytuł utrzymywał przez 41 rok, dopóki budynek Chrysler w Nowym Jorku nie został ukończony w 1930 roku."
     }
     ```
 
+#### Generation parameters
+
+The `generation_parameters` field in the request body allows you to specify the parameters which are described in the model documentation.
+
+[For Seamless model](https://huggingface.co/docs/transformers/main/en/model_doc/seamless_m4t#transformers.SeamlessM4TForTextToText.generate) and [for mBART model](https://huggingface.co/docs/transformers/main/en/model_doc/mbart#transformers.MBartForConditionalGeneration.generate)
+
 ### Health Check
 
 - Request:

docs/README.md

@@ -89,20 +89,32 @@ Choose your preferred distribution:
 - Request:
 
     ```bash
-    curl -X POST "http://localhost:8000/translate" \
-        -F "text=Hello, how are you?" \
-        -F "source_language=en_US"
-        -F "target_language=pl_PL"
+    curl -X 'POST' \
+      'http://127.0.0.1:8000/translate' \
+      -H 'accept: application/json' \
+      -H 'Content-Type: application/json' \
+      -d '{
+      "text_to_translate": "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct.",
+      "source_language": "en_US",
+      "target_language": "pl_PL",
+      "generation_parameters": { "max_length": 10240, "num_beams": 10 }
+    }'
     ```
 
 - Response:
 
     ```json
     {
-      "content": "Cześć, jak się masz?",
+      "translation": "Wieża Eiffla ma wysokość 324 metrów, mniej więcej taką samą wysokość jak 81-piętrowy budynek, i jest najwyższą budowlą w Paryżu. Jego podstawa jest kwadratowa, mierząc 125 metrów na każdej stronie. Podczas jej budowy Wieża Eiffla przekroczyła Pomnik Waszyngtonu, stając się najwyższą budowlą stworzoną przez człowieka na świecie, tytuł utrzymywał przez 41 rok, dopóki budynek Chrysler w Nowym Jorku nie został ukończony w 1930 roku."
     }
     ```
 
+#### Generation parameters
+
+The `generation_parameters` field in the request body allows you to specify the parameters which are described in the model documentation.
+
+[For Seamless model](https://huggingface.co/docs/transformers/main/en/model_doc/seamless_m4t#transformers.SeamlessM4TForTextToText.generate) and [for mBART model](https://huggingface.co/docs/transformers/main/en/model_doc/mbart#transformers.MBartForConditionalGeneration.generate)
+
 ### Health Check
 
 - Request:
@@ -156,6 +168,7 @@ Developer guide is available in [docs/DEVELOPER.md](DEVELOPER.md).
     - [Using Windows Executable](#using-windows-executable)
   - [API Features](#api-features)
     - [Translate Text](#translate-text)
+      - [Generation parameters](#generation-parameters)
     - [Health Check](#health-check)
   - [Configuration](#configuration)
   - [Supported Languages](#supported-languages)

src/api/dtos/translate_dto.py

@@ -1,15 +1,16 @@
 import re
-from typing import Optional
+from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, field_validator
 
 from domain.exceptions.invalid_language_format_error import InvalidLanguageFormatError
 
 
 class TranslateDTO(BaseModel):
-    text: str
+    text_to_translate: str
     source_language: str
     target_language: str
+    generation_parameters: Dict[str, Any] = {}
 
     @staticmethod
     def validate_language_format(v: str) -> str:

src/api/dtos/translate_result_dto.py

@@ -2,4 +2,4 @@
 
 
 class TranslateResultDTO(BaseModel):
-    content: str
+    translation: str

src/api/routers/translate_router.py

@@ -1,6 +1,6 @@
 from typing import Annotated
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Body, Depends
 
 from api.dtos.translate_dto import TranslateDTO
 from api.dtos.translate_result_dto import TranslateResultDTO
@@ -15,14 +15,15 @@ def __init__(self) -> None:
     async def translate(
         self,
         translate_text_usecase: Annotated[TranslateTextUseCase, Depends()],
-        translate_dto: TranslateDTO = Depends(),
+        translate_dto: TranslateDTO = Body(...),
     ) -> TranslateResultDTO:
-        result = await translate_text_usecase.execute(
-            translate_dto.text,
+        translation = await translate_text_usecase.execute(
+            translate_dto.text_to_translate,
             translate_dto.source_language,
             translate_dto.target_language,
+            translate_dto.generation_parameters,
         )
 
         return TranslateResultDTO(
-            content=result,
+            translation=translation,
         )

src/application/usecases/translate_text_usecase.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Any, Dict
 
 from fastapi import Depends
 
@@ -20,20 +20,22 @@ def __init__(
 
     async def execute(
         self,
-        text: str,
+        text_to_translate: str,
         source_language: str,
         target_language: str,
+        generation_parameters: Dict[str, Any],
     ) -> str:
         self.logger.info(
-            f"Executing translation for text '{text}' from '{source_language}' to '{target_language}'",
+            f"Executing translation for text '{text_to_translate}' from '{source_language}' to '{target_language}'",
         )
 
-        translation_result: str = self.translation_service.translate_text(
-            text,
+        translation: str = self.translation_service.translate_text(
+            text_to_translate,
             source_language,
             target_language,
+            generation_parameters,
         )
 
         self.logger.info("Returning translation result")
 
-        return translation_result
+        return translation

src/data/repositories/translation_model_repository_impl.py

@@ -1,6 +1,6 @@
 import threading
 import time
-from typing import Annotated, Optional
+from typing import Annotated, Any, Dict, Optional
 
 from fastapi import Depends
 
@@ -59,9 +59,10 @@ def _check_idle_timeout(self) -> None:
 
     def translate(
         self,
-        text: str,
+        text_to_translate: str,
         source_language: str,
         target_language: str,
+        generation_parameters: Dict[str, Any],
     ) -> str:
         with self._lock:
             if not self.worker.is_alive():
@@ -72,10 +73,11 @@ def translate(
             f"Translating started from source_language: {source_language}, target_language: {target_language}",
         )
 
-        result: str = self.worker.translate(
-            text,
+        translation: str = self.worker.translate(
+            text_to_translate,
             source_language,
             target_language,
+            generation_parameters,
         )
 
         self.timer.start(
@@ -89,4 +91,4 @@ def translate(
             f"Translating completed from source_language: {source_language}, target_language: {target_language}",
         )
 
-        return result
+        return translation

src/data/workers/mbart_translation_worker.py

@@ -3,9 +3,8 @@
 import multiprocessing.synchronize
 from dataclasses import dataclass
 from multiprocessing.sharedctypes import Synchronized
-from typing import Tuple
+from typing import Any, Dict, Tuple
 
-import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 from data.workers.base_worker import BaseWorker
@@ -22,22 +21,33 @@ class MBartTranslationConfig:
 
 class MBartTranslationWorker(
     BaseWorker[  # type: ignore
-        Tuple[str, str, str],
+        Tuple[str, str, str, Dict[str, Any]],
         str,
         MBartTranslationConfig,
         Tuple[AutoModelForSeq2SeqLM, AutoTokenizer],
     ],
 ):
     def translate(
         self,
-        text: str,
+        text_to_translate: str,
         source_language: str,
         target_language: str,
+        generation_parameters: Dict[str, Any],
     ) -> str:
         if not self.is_alive():
             raise WorkerNotRunningError()
 
-        self._pipe_parent.send(("translate", (text, source_language, target_language)))
+        self._pipe_parent.send(
+            (
+                "translate",
+                (
+                    text_to_translate,
+                    source_language,
+                    target_language,
+                    generation_parameters,
+                ),
+            ),
+        )
         result = self._pipe_parent.recv()
 
         if isinstance(result, Exception):
@@ -62,7 +72,7 @@ def initialize_shared_object(
     def handle_command(
         self,
         command: str,
-        args: Tuple[str, str, str],
+        args: Tuple[str, str, str, Dict[str, Any]],
         shared_object: Tuple[AutoModelForSeq2SeqLM, AutoTokenizer],
         config: MBartTranslationConfig,
         pipe: multiprocessing.connection.Connection,
@@ -74,21 +84,20 @@ def handle_command(
                 with processing_lock:
                     is_processing.value = True
 
-                text, source_language, target_language = args
+                text, source_language, target_language, generation_parameters = args
                 model, tokenizer = shared_object
 
                 tokenizer.src_lang = source_language
-                inputs = tokenizer([text], truncation=True, padding=True, max_length=1024, return_tensors="pt")
-
-                for key in inputs:
-                    inputs[key] = inputs[key].to(config.device)
+                inputs = tokenizer(text, return_tensors="pt").to(config.device)
 
-                with torch.no_grad():
+                if "forced_bos_token_id" in generation_parameters:
+                    kwargs = {"forced_bos_token_id": generation_parameters["forced_bos_token_id"]}
+                else:
                     kwargs = {"forced_bos_token_id": tokenizer.lang_code_to_id[target_language]}
 
-                    translated = model.generate(**inputs, num_beams=5, **kwargs)
+                translation = model.generate(**inputs, **kwargs)
 
-                    output = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
+                output = [tokenizer.decode(t, skip_special_tokens=True) for t in translation]
 
                 pipe.send("".join(output))
 

src/data/workers/seamless_translation_worker.py

@@ -3,7 +3,7 @@
 import multiprocessing.synchronize
 from dataclasses import dataclass
 from multiprocessing.sharedctypes import Synchronized
-from typing import Tuple
+from typing import Any, Dict, Tuple
 
 from transformers import AutoProcessor, SeamlessM4Tv2ForTextToText
 
@@ -21,22 +21,33 @@ class SeamlessTranslationConfig:
 
 class SeamlessTranslationWorker(
     BaseWorker[  # type: ignore
-        Tuple[str, str, str],
+        Tuple[str, str, str, Dict[str, Any]],
         str,
         SeamlessTranslationConfig,
         Tuple[SeamlessM4Tv2ForTextToText, AutoProcessor],
     ],
 ):
     def translate(
         self,
-        text: str,
+        text_to_translate: str,
         source_language: str,
         target_language: str,
+        generation_parameters: Dict[str, Any],
     ) -> str:
         if not self.is_alive():
             raise WorkerNotRunningError()
 
-        self._pipe_parent.send(("translate", (text, source_language, target_language)))
+        self._pipe_parent.send(
+            (
+                "translate",
+                (
+                    text_to_translate,
+                    source_language,
+                    target_language,
+                    generation_parameters,
+                ),
+            ),
+        )
         result = self._pipe_parent.recv()
 
         if isinstance(result, Exception):
@@ -61,7 +72,7 @@ def initialize_shared_object(
     def handle_command(
         self,
         command: str,
-        args: Tuple[str, str, str],
+        args: Tuple[str, str, str, Dict[str, Any]],
         shared_object: Tuple[SeamlessM4Tv2ForTextToText, AutoProcessor],
         config: SeamlessTranslationConfig,
         pipe: multiprocessing.connection.Connection,
@@ -73,15 +84,26 @@ def handle_command(
                 with processing_lock:
                     is_processing.value = True
 
-                text, source_language, target_language = args
+                text, source_language, target_language, generation_parameters = args
                 model, processor = shared_object
 
-                processor.src_lang = source_language
-                input_tokens = processor(text, return_tensors="pt", padding=True).to(config.device)
-
-                output_tokens = model.generate(**input_tokens, tgt_lang=target_language)[0].tolist()
-
-                text_output = processor.decode(output_tokens, skip_special_tokens=True)
+                input_tokens = processor(
+                    text,
+                    src_lang=source_language,
+                    return_tensors="pt",
+                    padding=True,
+                ).to(config.device)
+
+                output_tokens = model.generate(
+                    **input_tokens,
+                    tgt_lang=target_language,
+                    **generation_parameters,
+                )[0].tolist()
+
+                text_output = processor.decode(
+                    output_tokens,
+                    skip_special_tokens=True,
+                )
 
                 pipe.send("".join(text_output))