✨ Allow setting model / language on new doc page

Fixes #264 Fixes #208
bugbakery · Jul 6, 2023 · 7a1a3c4 · 7a1a3c4
1 parent 28752bc
commit 7a1a3c4
Show file tree

Hide file tree

Showing 12 changed files with 510 additions and 12 deletions.
diff --git a/backend/data/models.json b/backend/data/models.json
@@ -0,0 +1,205 @@
+{
+    "tiny.en": {
+        "id": "tiny.en",
+        "name": "Tiny En",
+        "languages": [
+            "auto"
+        ]
+    },
+    "tiny": {
+        "id": "tiny",
+        "name": "Tiny",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    },
+    "base.en": {
+        "id": "base.en",
+        "name": "Base En",
+        "languages": [
+            "auto"
+        ]
+    },
+    "base": {
+        "id": "base",
+        "name": "Base",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    },
+    "small.en": {
+        "id": "small.en",
+        "name": "Small En",
+        "languages": [
+            "auto"
+        ]
+    },
+    "small.en-tdrz": {
+        "id": "small.en-tdrz",
+        "name": "Small En Tdrz",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    },
+    "small": {
+        "id": "small",
+        "name": "Small",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    },
+    "medium.en": {
+        "id": "medium.en",
+        "name": "Medium En",
+        "languages": [
+            "auto"
+        ]
+    },
+    "medium": {
+        "id": "medium",
+        "name": "Medium",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    },
+    "large-v1": {
+        "id": "large-v1",
+        "name": "Large V1",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    },
+    "large": {
+        "id": "large",
+        "name": "Large",
+        "languages": [
+            "auto",
+            "ar",
+            "de",
+            "el",
+            "en",
+            "es",
+            "fa",
+            "fi",
+            "fr",
+            "hu",
+            "it",
+            "ja",
+            "nl",
+            "pl",
+            "pt",
+            "ru",
+            "tr",
+            "uk",
+            "zh"
+        ]
+    }
+}
diff --git a/backend/openapi-schema.yml b/backend/openapi-schema.yml
@@ -79,11 +79,19 @@ components:
           format: binary
           title: File
           type: string
+        language:
+          title: Language
+          type: string
+        model:
+          title: Model
+          type: string
         name:
           title: Name
           type: string
       required:
       - name
+      - model
+      - language
       - file
       title: Body_create_document_api_v1_documents__post
       type: object
@@ -192,6 +200,36 @@ components:
       - token
       title: LoginResponse
       type: object
+    ModelConfig:
+      properties:
+        id:
+          title: Id
+          type: string
+        languages:
+          items:
+            type: string
+          title: Languages
+          type: array
+        name:
+          title: Name
+          type: string
+      required:
+      - id
+      - name
+      - languages
+      title: ModelConfig
+      type: object
+    PublicConfig:
+      properties:
+        models:
+          additionalProperties:
+            $ref: '#/components/schemas/ModelConfig'
+          title: Models
+          type: object
+      required:
+      - models
+      title: PublicConfig
+      type: object
     SetDurationRequest:
       properties:
         duration:
@@ -365,6 +403,17 @@ paths:
               schema: {}
           description: Successful Response
       summary: Root
+  /api/v1/config/:
+    get:
+      operationId: get_config_api_v1_config__get
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/PublicConfig'
+          description: Successful Response
+      summary: Get Config
   /api/v1/documents/:
     get:
       operationId: list_documents_api_v1_documents__get

diff --git a/backend/tests/test_doc.py b/backend/tests/test_doc.py
@@ -16,7 +16,9 @@
 @pytest.fixture
 def document(memory_session: Session, logged_in_client: TestClient):
     req = logged_in_client.post(
-        "/api/v1/documents/", files={"file": b""}, data={"name": "test document"}
+        "/api/v1/documents/",
+        files={"file": b""},
+        data={"name": "test document", "model": "tiny", "language": "auto"},
     )
     assert req.status_code == 200
     document_id = req.json()["id"]
@@ -48,7 +50,9 @@ def test_doc_delete(
     files = set(str(x) for x in settings.storage_path.glob("*"))
 
     req = logged_in_client.post(
-        "/api/v1/documents/", files={"file": b""}, data={"name": "test document"}
+        "/api/v1/documents/",
+        files={"file": b""},
+        data={"name": "test document", "model": "tiny", "language": "auto"},
     )
     assert req.status_code == 200
     document_id = req.json()["id"]

diff --git a/backend/tests/test_performance.py b/backend/tests/test_performance.py
@@ -7,7 +7,9 @@
 
 def create_doc(logged_in_client: TestClient):
     req = logged_in_client.post(
-        "/api/v1/documents/", files={"file": b""}, data={"name": "test document"}
+        "/api/v1/documents/",
+        files={"file": b""},
+        data={"name": "test document", "model": "tiny", "language": "auto"},
     )
     assert req.status_code == 200
 

diff --git a/backend/transcribee_backend/config.py b/backend/transcribee_backend/config.py
@@ -1,6 +1,8 @@
 from pathlib import Path
+from typing import Dict, List
 
-from pydantic import BaseSettings
+from pydantic import BaseSettings, parse_file_as
+from pydantic.main import BaseModel
 
 
 class Settings(BaseSettings):
@@ -12,5 +14,25 @@ class Settings(BaseSettings):
 
     media_url_base = "http://localhost:8000/"
 
+    model_config_path: Path = Path("data/models.json")
+
+
+class ModelConfig(BaseModel):
+    id: str
+    name: str
+    languages: List[str]
+
+
+class PublicConfig(BaseModel):
+    models: Dict[str, ModelConfig]
+
+
+def get_model_config():
+    return parse_file_as(Dict[str, ModelConfig], settings.model_config_path)
+
+
+def get_public_config():
+    return PublicConfig(models=get_model_config())
+
 
 settings = Settings()
diff --git a/backend/transcribee_backend/main.py b/backend/transcribee_backend/main.py
@@ -6,6 +6,7 @@
 from transcribee_backend.config import settings
 from transcribee_backend.helpers.periodic_tasks import run_periodic
 from transcribee_backend.helpers.tasks import timeout_attempts
+from transcribee_backend.routers.config import config_router
 from transcribee_backend.routers.document import document_router
 from transcribee_backend.routers.task import task_router
 from transcribee_backend.routers.user import user_router
@@ -28,6 +29,7 @@
 app.include_router(user_router, prefix="/api/v1/users")
 app.include_router(document_router, prefix="/api/v1/documents")
 app.include_router(task_router, prefix="/api/v1/tasks")
+app.include_router(config_router, prefix="/api/v1/config")
 
 
 @app.get("/")

diff --git a/backend/transcribee_backend/routers/config.py b/backend/transcribee_backend/routers/config.py
@@ -0,0 +1,10 @@
+from fastapi import APIRouter
+
+from transcribee_backend.config import PublicConfig, get_public_config
+
+config_router = APIRouter()
+
+
+@config_router.get("/")
+def get_config() -> PublicConfig:
+    return get_public_config()