Added model config and dynamic pipeline loading

.gitignore

@@ -0,0 +1,166 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python
+# Edit at https://www.toptal.com/developers/gitignore?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# End of https://www.toptal.com/developers/gitignore/api/python

Untitled.ipynb

@@ -0,0 +1,162 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "id": "9654bdb9-79ea-49b0-ab6c-da9ddbdb7ee2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.utils import load_pipelines_from_config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "b3b42ff6-dea6-4700-ab96-6bf1706221b7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pipelines = load_pipelines_from_config(\"model_config.yaml\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "667ae085-587a-427a-87dd-bfc7c08553a5",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'Rap': <transformers.pipelines.text_generation.TextGenerationPipeline at 0x7f5650745420>,\n",
+       " 'Metal': <transformers.pipelines.text_generation.TextGenerationPipeline at 0x7f5595d74700>}"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pipelines"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "0a90af2f-2f8f-4c0a-a2ef-ce8d7bf5e4be",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"model_config.yaml\", \"r\") as f:\n",
+    "    model_config = yaml.safe_load(f.read())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "e743aa51-6d6b-43c6-b138-d9b7e9b70a36",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'Rap': {'model_name': 'nschenone/rap-distil',\n",
+       "  'artist_names': ['Eminem', 'Hopsin', 'Kentrick Lamar'],\n",
+       "  'mlflow_run_id': '16c4ff05d92a45d79d89572a58b6424b',\n",
+       "  'hf_commit_hash': 'ca066f322213fbeac8d036fafd32112e23837722',\n",
+       "  'task': 'text-generation'},\n",
+       " 'Metal': {'model_name': 'nschenone/metal-distil',\n",
+       "  'artist_names': ['Slipknot', 'Parkway Drive', 'Periphery'],\n",
+       "  'mlflow_run_id': 'f30f57e3d8c440a09e1738f07db0b211',\n",
+       "  'hf_commit_hash': 'ed0657933ac3eb11a554dbe153363ff3e457f5ab',\n",
+       "  'task': 'text-generation'}}"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "8608031c-c0a2-44ac-95fb-3acfde35c125",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "a25764a5-e0db-45a5-96ed-34efde27046a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "a = pipeline(\n",
+    "    task=\"text-generation\",\n",
+    "    model=\"nschenone/rap-distil\",\n",
+    "    revision=\"753f2768e0a9d5b21b5009bec4855ed2c2ddef16\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "632a4847-75e4-459b-a4aa-248ff741f879",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "ValueError",
+     "evalue": "You need to specify a `repo_path_or_name` or a `repo_url`.",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Input \u001b[0;32mIn [23]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43ma\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpush_to_hub\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
+      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py:2620\u001b[0m, in \u001b[0;36mPreTrainedModel.push_to_hub\u001b[0;34m(self, repo_path_or_name, repo_url, use_temp_dir, commit_message, organization, private, use_auth_token, max_shard_size, **model_card_kwargs)\u001b[0m\n\u001b[1;32m   2617\u001b[0m     repo_path_or_name \u001b[38;5;241m=\u001b[39m tempfile\u001b[38;5;241m.\u001b[39mmkdtemp()\n\u001b[1;32m   2619\u001b[0m \u001b[38;5;66;03m# Create or clone the repo. If the repo is already cloned, this just retrieves the path to the repo.\u001b[39;00m\n\u001b[0;32m-> 2620\u001b[0m repo \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_create_or_get_repo\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   2621\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrepo_path_or_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrepo_path_or_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2622\u001b[0m \u001b[43m    \u001b[49m\u001b[43mrepo_url\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mrepo_url\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2623\u001b[0m \u001b[43m    \u001b[49m\u001b[43morganization\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43morganization\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2624\u001b[0m \u001b[43m    \u001b[49m\u001b[43mprivate\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mprivate\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2625\u001b[0m \u001b[43m    \u001b[49m\u001b[43muse_auth_token\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43muse_auth_token\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   2626\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   2627\u001b[0m \u001b[38;5;66;03m# Save the files in the cloned repo\u001b[39;00m\n\u001b[1;32m   2628\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msave_pretrained(repo_path_or_name, max_shard_size\u001b[38;5;241m=\u001b[39mmax_shard_size)\n",
+      "File \u001b[0;32m/opt/conda/lib/python3.10/site-packages/transformers/utils/hub.py:1017\u001b[0m, in \u001b[0;36mPushToHubMixin._create_or_get_repo\u001b[0;34m(cls, repo_path_or_name, repo_url, organization, private, use_auth_token)\u001b[0m\n\u001b[1;32m   1007\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m   1008\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_create_or_get_repo\u001b[39m(\n\u001b[1;32m   1009\u001b[0m     \u001b[38;5;28mcls\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1014\u001b[0m     use_auth_token: Optional[Union[\u001b[38;5;28mbool\u001b[39m, \u001b[38;5;28mstr\u001b[39m]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m   1015\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Repository:\n\u001b[1;32m   1016\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m repo_path_or_name \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m repo_url \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m-> 1017\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou need to specify a `repo_path_or_name` or a `repo_url`.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m   1019\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m use_auth_token \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m repo_url \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m   1020\u001b[0m         use_auth_token \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
+      "\u001b[0;31mValueError\u001b[0m: You need to specify a `repo_path_or_name` or a `repo_url`."
+     ]
+    }
+   ],
+   "source": [
+    "a.model.push_to_hub()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2a0c4d95-8ea7-4b40-98f1-3ad16c31eeba",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

app.py

@@ -1,16 +1,9 @@
 import gradio as gr
-from transformers import pipeline, set_seed
+from transformers import set_seed
 
-models = {
-    "Rap" : pipeline(
-        task="text-generation",
-        model="nschenone/rap-distil"
-    ),
-    "Metal" : pipeline(
-        task="text-generation",
-        model="nschenone/metal-distil"
-    )
-}
+from src.utils import load_pipelines_from_config
+
+pipelines = load_pipelines_from_config(config_path="model_config.yaml")
 
 def generate(
     text: str,
@@ -27,7 +20,7 @@ def generate(
 
     set_seed(seed)
 
-    generated = models[model](
+    generated = pipelines[model](
         text_inputs=text,
         max_length=max_length,
         num_return_sequences=num_return_sequences,
@@ -50,8 +43,8 @@ iface = gr.Interface(
             label="Input Text"
         ),
         gr.Dropdown(
-            choices=list(models.keys()),
-            value=list(models.keys())[0],
+            choices=list(pipelines.keys()),
+            value=list(pipelines.keys())[0],
             label="Model"
         ),
         gr.Slider(

model_config.yaml

@@ -0,0 +1,18 @@
+Rap:
+  model_name: nschenone/rap-distil
+  artist_names: 
+  - Eminem
+  - Hopsin
+  - Kentrick Lamar
+  mlflow_run_id: 16c4ff05d92a45d79d89572a58b6424b
+  hf_commit_hash: ca066f322213fbeac8d036fafd32112e23837722
+  task: text-generation
+Metal:
+  model_name: nschenone/metal-distil
+  artist_names:
+  - Slipknot
+  - Parkway Drive
+  - Periphery
+  mlflow_run_id: f30f57e3d8c440a09e1738f07db0b211
+  hf_commit_hash: ed0657933ac3eb11a554dbe153363ff3e457f5ab
+  task: text-generation

src/utils.py

@@ -0,0 +1,16 @@
+import yaml
+from transformers import pipeline
+
+def load_pipelines_from_config(config_path: str):
+    with open(config_path, "r") as f:
+        model_config = yaml.safe_load(f.read())
+        
+    models = {}
+    for model, config in model_config.items():
+        models[model] = pipeline(
+            task=config["task"],
+            model=config["model_name"],
+            revision=config["hf_commit_hash"]
+        )
+        
+    return models