feat: Support MedCAT v2 (#25)

* feat: support MedCAT v2
* feat: make MedCAT V2 ontology mappings configurable
* fix: add workaround for metrics collection for supervsised training
* feat: add concept ids to evaluation results and deprecate py38 support
* chore: update individual local dev containers
* chore: improve type hints
* feat: support python 3.12

Author: baixiac

.github/workflows/main.yaml

@@ -12,19 +12,25 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.8', '3.9', '3.10', '3.11' ]
+        python-version: [
+          '3.9',
+          '3.10',
+          '3.11',
+          '3.12',
+        ]
       max-parallel: 4
 
     steps:
       - uses: actions/checkout@v4
       - name: Install uv and set Python to ${{ matrix.python-version }}
         uses: astral-sh/setup-uv@v6
         with:
-          version: "0.7.20"
+          version: "0.8.10"
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          uv sync --group dev --group docs
+          uv sync --extra dev --extra docs --extra vllm
+          uv run python -m ensurepip
       - name: Check types
         run: |
           uv run mypy app

.github/workflows/release.yaml

@@ -17,11 +17,11 @@ jobs:
       - name: Install uv
         uses: astral-sh/setup-uv@v5
         with:
-            version: "0.6.10"
+            version: "0.8.10"
             python-version: "3.10"
       - name: Install dependencies
         run: |
-          uv sync --group dev --group docs --group vllm
+          uv sync --extra dev --extra docs --extra vllm
       - name: Run unit tests
         run: |
           uv run pytest -v tests/app --cov --cov-report=html:coverage_reports #--random-order

app/api/api.py

@@ -211,9 +211,9 @@ def _get_app(
         tags = TagsStreamable
     else:
         tags = Tags
-    tags_metadata = [{  # type: ignore
-        "name": tag.name,   # type: ignore
-        "description": tag.value    # type: ignore
+    tags_metadata = [{
+        "name": tag.name,
+        "description": tag.value
     } for tag in tags]
     app = FastAPI(
         title="CogStack ModelServe",

app/api/auth/db.py

@@ -52,5 +52,4 @@ async def get_user_db(session: AsyncSession = Depends(_get_async_session)) -> As
         SQLAlchemyUserDatabase: A database instance initialised with the given session and the User model.
     """
 
-    # TODO: fix this type checking error
     yield SQLAlchemyUserDatabase(session, User)

app/api/routers/generative.py

@@ -217,7 +217,8 @@ async def _stream(prompt: str, max_tokens: int, temperature: float) -> AsyncGene
             yield f"data: {json.dumps(data)}\n\n"
         yield "data: [DONE]\n\n"
 
-    prompt = get_prompt_from_messages(model_service.tokenizer, messages)    # type: ignore
+    assert hasattr(model_service, "tokenizer"), "Model service doesn't have a tokenizer"
+    prompt = get_prompt_from_messages(model_service.tokenizer, messages)
     if stream:
         return StreamingResponse(
             _stream(prompt, max_tokens, temperature),

app/api/routers/stream.py

@@ -137,7 +137,7 @@ def __init__(
         self,
         content: Any,
         status_code: int = 200,
-        max_chunk_size: Optional[int] = 1024,
+        max_chunk_size: int = 1024,
         headers: Optional[Mapping[str, str]] = None,
         media_type: Optional[str] = None,
         background: Optional[BackgroundTask] = None,
@@ -161,8 +161,8 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                 })
                 response_started = True
             line_bytes = line.encode("utf-8")
-            for i in range(0, len(line_bytes), self.max_chunk_size):    # type: ignore
-                chunk = line_bytes[i:i + self.max_chunk_size]   # type: ignore
+            for i in range(0, len(line_bytes), self.max_chunk_size):
+                chunk = line_bytes[i:i + self.max_chunk_size]
                 await send({
                     "type": "http.response.body",
                     "body": chunk,

app/api/utils.py

@@ -367,7 +367,7 @@ async def generate_text(
         params = SamplingParams(max_tokens=max_tokens)
 
         conversation, _ = parse_chat_messages(messages, model_config, tokenizer, content_format="string")   # type: ignore
-        prompt_tokens = apply_hf_chat_template(    # type: ignore
+        prompt_tokens = apply_hf_chat_template( # type: ignore
             tokenizer,
             conversation=conversation,
             tools=None,

app/config.py

@@ -36,6 +36,7 @@ class Settings(BaseSettings):   # type: ignore
     TRAINING_CACHE_DIR: str = os.path.join(os.path.abspath(os.path.dirname(__file__)), "cms_cache")           # the directory to cache the intermediate files created during training
     HF_PIPELINE_AGGREGATION_STRATEGY: str = "simple"  # the strategy used for aggregating the predictions of the Hugging Face NER model
     LOG_PER_CONCEPT_ACCURACIES: str = "false"         # if "true", per-concept accuracies will be exposed to the metrics scrapper. Switch this on with caution due to the potentially high number of concepts
+    MEDCAT2_MAPPED_ONTOLOGIES: str = ""               # the comma-separated names of ontologies for MedCAT2 to map to
     DEBUG: str = "false"                              # if "true", the debug mode is switched on
 
     class Config:

app/envs/.env

@@ -73,5 +73,8 @@ TRAINING_SAFE_MODEL_SERIALISATION=false
 # The strategy used for aggregating the predictions of the Hugging Face NER model
 HF_PIPELINE_AGGREGATION_STRATEGY=simple
 
+# The comma-separated names of ontologies for MedCAT2 to map to
+MEDCAT2_MAPPED_ONTOLOGIES=opcs4,icd10
+
 # If "true", the debug mode is switched on
 DEBUG=false

app/management/tracker_client.py

@@ -1,3 +1,4 @@
+import ast
 import os
 import socket
 import mlflow
@@ -114,7 +115,7 @@ def send_model_stats(stats: Dict, step: int) -> None:
             step (int): The current step in the training or evaluation process.
         """
 
-        metrics = {key.replace(" ", "_").lower(): val for key, val in stats.items()}
+        metrics = {key.replace(" ", "_").lower(): val for key, val in stats.items() if isinstance(val, (int, float))}
         mlflow.log_metrics(metrics, step)
 
     @staticmethod
@@ -563,6 +564,7 @@ def get_metrics_by_job_id(self, job_id: str) -> List[Dict[str, Any]]:
                 metrics_history = {}
                 for metric in run.data.metrics.keys():
                     metrics_history[metric] = [m.value for m in self.mlflow_client.get_metric_history(run_id=run.info.run_id, key=metric)]
+                metrics_history["concepts"] = ast.literal_eval(run.data.tags.get("training.entity.classes", "[]"))
                 metrics.append(metrics_history)
             return metrics
         except MlflowException as e: