Merge branch 'ml4ai:main' into main

Author: bigglesandginger

.drone.yml

@@ -81,14 +81,17 @@ jobs:
       working-directory: .
       run: |
         # retrieve latest model for img2mml component
-        curl -L https://artifacts.askem.lum.ai/skema/img2mml/models/cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt > skema/img2mml/trained_models/cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt
+        pip install huggingface_hub
+        python scripts/retrieve_model_ci.py
+
+        # Install askem
         pip install ".[all]"
+
     # Install tree-sitter parser (for Python component unit tests)
     - name: Install tree-sitter parsers
       working-directory: .
       run: python skema/program_analysis/tree_sitter_parsers/build_parsers.py --ci --all
 
-
     # docs (API)
     # generate python docs using pdoc
     - name: "Create documentation for Python components (API docs)"

.github/workflows/tests-and-docs.yml

@@ -54,7 +54,7 @@ ENV PATH="/root/.cargo/bin:${PATH}"
 RUN pip install wheel
 RUN pip install six
 # Download ML model (~150MB)
-RUN curl -L https://artifacts.askem.lum.ai/skema/img2mml/models/cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt > skema/img2mml/trained_models/cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt
+RUN pip install huggingface_hub && python scripts/retrieve_model_ci.py
 RUN tree /app
 #RUN pip install ".[all]"
 # exclude dependencies for docs

Dockerfile.skema-py

@@ -9,7 +9,7 @@ conda activate skema
 # Install tree-sitter parsers
 python skema/program_analysis/tree_sitter_parsers/build_parsers.py --all
 # download the checkpoint for the img2mml service
-curl -L https://artifacts.askem.lum.ai/skema/img2mml/models/cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt > skema/img2mml/trained_models/cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt
+python scripts/retrieve_model.py
 # mathjax deps for img2mml
 (cd skema/img2mml/data_generation && npm install)
 ```

docs/dev/env.md

@@ -15,19 +15,19 @@ dependencies=[
     "numpy",
     "dill==0.3.7",
     "networkx==2.8.8",
-    "PyYAML",
+    "PyYAML==6.*",
     "tree-sitter==0.20.4",
     "neo4j==5.14.1",
     "requests",
-    "beautifulsoup4", # used to remove comments etc from pMML before sending to MORAE
+    "beautifulsoup4==4.12.*", # used to remove comments etc from pMML before sending to MORAE
     "typing_extensions", # see https://github.com/pydantic/pydantic/issues/5821#issuecomment-1559196859
     "fastapi~=0.100.0",
     "starlette",
     "httpx",
     "pydantic>=2.0.0",
     "uvicorn",
     "python-multipart",
-    "func_timeout"
+    "func_timeout==4.3.5"
 ]
 # The Python program analysis pipeline does not currently work with Python 3.9
 # or 3.10. This may change in the future.
@@ -57,7 +57,7 @@ isa = [
 ]
 
 # shared ML dependencies
-ml = ["torch==2.0.1", "torchvision==0.15.2", "beartype==0.15.0"]
+ml = ["torch==2.0.1", "torchvision==0.15.2", "beartype==0.15.0", "huggingface_hub"]
 
 # Im2MML dependencies. The img2mml service converts equation images to MathML.
 # See the skema/img2mml directory.
@@ -100,6 +100,7 @@ all = ["skema[core]", "skema[dev]", "skema[doc]", "skema[demo]", "skema[annotati
 "skema.rest" = "skema/rest"
 "skema.skema_py" = "skema/skema_py"
 "skema.utils" = "skema/utils"
+"skema.data" = "skema/data"
 
 # re-map skema/text_reading/python to skema.text_reading
 #"skema.text_reading" = "skema/text_reading/python"
@@ -110,7 +111,7 @@ all = ["skema[core]", "skema[dev]", "skema[doc]", "skema[demo]", "skema[annotati
 
 [tool.setuptools.package-data]
 # needed to ensure models are included in package/discoverable
-"*" = ["*.json", "vocab.txt", "*.pt", "*.png", "*.html", "*.yml", "*.yaml"]
+"*" = ["*.json", "*vocab.txt", "*.pt", "*.png", "*.html", "*.yml", "*.yaml", "*.zip"]
 
 [tool.setuptools.dynamic]
 readme = {file = ["README.md"], content-type = "text/markdown"}

pyproject.toml

@@ -0,0 +1,31 @@
+import os
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+
+def retrieve_model(model_path=None) -> str:
+    """
+    Retrieve the img2mml model from the specified path or download it if not found.
+
+    Args:
+        model_path (str, optional): Path to the img2mml model file. Defaults to None.
+
+    Returns:
+        str: Path to the loaded model file.
+    """
+    cwd = Path(__file__).parents[0]
+    REPO_NAME = "lum-ai/img2mml"
+    MODEL_NAME = "cnn_xfmer_arxiv_im2mml_with_fonts_boldface_best.pt"
+    # If the model path is none or doesn't exist, the default model will be downloaded from server.
+    if model_path is None or not os.path.exists(model_path):
+        model_path = cwd / "trained_models" / MODEL_NAME
+
+        # Check if the model file already exists
+        if not os.path.exists(model_path):
+            # If the file doesn't exist, download it from the specified URL
+            print(f"Downloading the model checkpoint from HuggingFace...")
+            hf_hub_download(repo_id=REPO_NAME, filename=MODEL_NAME, local_dir=model_path.parent, local_dir_use_symlinks=False)
+        
+    return str(model_path)
+
+retrieve_model()