backend: Introduce models endpoint

Introduce a new `/models` endpoint that returns a list of suggested
models for a given task.

Currently, the only supported task is that of summarization.

Refs #381

Signed-off-by: Dimitris Poulopoulos <[email protected]>

lumigator/python/mzai/backend/backend/api/router.py

@@ -1,6 +1,6 @@
 from fastapi import APIRouter
 
-from backend.api.routes import completions, datasets, health, jobs
+from backend.api.routes import completions, datasets, health, jobs, models
 from backend.api.tags import Tags
 
 API_V1_PREFIX = "/api/v1"
@@ -11,3 +11,4 @@
 api_router.include_router(jobs.router, prefix="/jobs", tags=[Tags.JOBS])
 api_router.include_router(jobs.router, prefix="/experiments", tags=[Tags.EXPERIMENTS])
 api_router.include_router(completions.router, prefix="/completions", tags=[Tags.COMPLETIONS])
+api_router.include_router(models.router, prefix="/models", tags=[Tags.MODELS])

lumigator/python/mzai/backend/backend/api/routes/models.py

@@ -0,0 +1,81 @@
+from fastapi import APIRouter
+from lumigator_schemas.extras import ListingResponse
+
+SUGGESTED_MODELS = [
+    {
+        "name:": "facebook/bart-large-cnn",
+        "uri": "hf://facebook/bart-large-cnn",
+        "description": "BART is a large-sized model fine-tuned on the CNN Daily Mail dataset.",
+    },
+    {
+        "name": "mikeadimech/longformer-qmsum-meeting-summarization",
+        "uri": "hf://mikeadimech/longformer-qmsum-meeting-summarization",
+        "description": (
+            "Longformer is a transformer model that is capable of processing long sequences."
+        ),
+    },
+    {
+        "name": "mrm8488/t5-base-finetuned-summarize-news",
+        "uri": "hf://mrm8488/t5-base-finetuned-summarize-news",
+        "description": (
+            "Google's T5 base fine-tuned on News Summary dataset for summarization downstream task."
+        ),
+    },
+    {
+        "name": "Falconsai/text_summarization",
+        "uri": "hf://Falconsai/text_summarization",
+        "description": (
+            "A fine-tuned variant of the T5 transformer model, designed for the task "
+            "of text summarization."
+        ),
+    },
+    {
+        "name": "mistralai/Mistral-7B-Instruct-v0.3",
+        "uri": "hf://mistralai/Mistral-7B-Instruct-v0.3",
+        "description": (
+            "Mistral-7B-Instruct-v0.3 is an instruct fine-tuned version of the Mistral-7B-v0.3."
+        ),
+    },
+    {
+        "name": "gpt-4o-mini",
+        "uri": "oai://gpt-4o-mini",
+        "description": "OpenAI's GPT-4o-mini model.",
+    },
+    {
+        "name": "gpt-4-turbo",
+        "uri": "oai://gpt-4-turbo",
+        "description": "OpenAI's GPT-4 Turbo model.",
+    },
+    {
+        "name": "open-mistral-7b",
+        "uri": "mistral://open-mistral-7b",
+        "description": "Mistral's 7B model.",
+    },
+    {
+        "name": "mistralai/Mistral-7B-Instruct-v0.2",
+        "uri": "llamafile://mistralai/Mistral-7B-Instruct-v0.2",
+        "description": "A llamafile package of Mistral's 7B Instruct model.",
+    },
+]
+
+router = APIRouter()
+
+
+@router.get("/{task_name}/")
+def get_suggested_models(tast_name: str) -> ListingResponse[dict]:
+    """Get a list of suggested models for the given task.
+
+    Args:
+        task_name (str): The task name.
+
+    Returns:
+        ListingResponse[str]: A list of suggested models.
+    """
+    if tast_name != "summarization":
+        return ListingResponse[dict].model_validate({"total": 0, "items": []})
+
+    return_data = {
+        "total": len(SUGGESTED_MODELS),
+        "items": SUGGESTED_MODELS,
+    }
+    return ListingResponse[dict].model_validate(return_data)

lumigator/python/mzai/backend/backend/api/tags.py

@@ -7,6 +7,7 @@ class Tags(str, Enum):
     JOBS = "jobs"
     COMPLETIONS = "completions"
     EXPERIMENTS = "experiments"
+    MODELS = "models"
 
 
 TAGS_METADATA = [
@@ -30,6 +31,10 @@ class Tags(str, Enum):
         "name": Tags.COMPLETIONS,
         "description": "Access models via external vendor endpoints",
     },
+    {
+        "name": Tags.MODELS,
+        "description": "Return a list of suggested models for a given task.",
+    },
 ]
 """Metadata to associate with route tags in the OpenAPI documentation.