backend: Move suggested models in a YAML file

Move the suggested models list in a YAML file. Read the YAML file during runtime to respond appropriately. Refs #381 Signed-off-by: Dimitris Poulopoulos <[email protected]>
mozilla-ai · Nov 19, 2024 · b7254b1 · b7254b1
1 parent 9eb69e7
commit b7254b1
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 58 deletions.
diff --git a/lumigator/python/mzai/backend/backend/api/routes/models.py b/lumigator/python/mzai/backend/backend/api/routes/models.py
@@ -1,62 +1,10 @@
+from pathlib import Path
+
+import yaml
 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.",
-    },
-]
+MODELS_CONFIG_PATH = Path(__file__).resolve().parent.parent.parent / "models_config.yaml"
 
 router = APIRouter()
 
@@ -74,8 +22,11 @@ def get_suggested_models(tast_name: str) -> ListingResponse[dict]:
     if tast_name != "summarization":
         return ListingResponse[dict].model_validate({"total": 0, "items": []})
 
+    with Path(MODELS_CONFIG_PATH).open() as file:
+        data = yaml.safe_load(file)
+
     return_data = {
-        "total": len(SUGGESTED_MODELS),
-        "items": SUGGESTED_MODELS,
+        "total": len(data),
+        "items": data,
     }
     return ListingResponse[dict].model_validate(return_data)
diff --git a/lumigator/python/mzai/backend/backend/models_config.yaml b/lumigator/python/mzai/backend/backend/models_config.yaml
@@ -0,0 +1,35 @@
+- 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."