Initial: HumanSignal huggingface_ner example patched for HF Spaces

.dockerignore

@@ -0,0 +1,18 @@
+# Exclude everything
+**
+
+# Include Dockerfile and docker-compose for reference (optional, decide based on your use case)
+!Dockerfile
+!docker-compose.yml
+
+# Include Python application files
+!*.py
+
+# Include requirements files
+!requirements*.txt
+
+# Include script
+!*.sh
+
+# Exclude specific requirements if necessary
+# requirements-test.txt (Uncomment if you decide to exclude this)

Dockerfile

@@ -0,0 +1,48 @@
+# syntax=docker/dockerfile:1
+ARG PYTHON_VERSION=3.11
+
+FROM python:${PYTHON_VERSION}-slim AS python-base
+ARG TEST_ENV
+
+WORKDIR /app
+
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PORT=${PORT:-9090} \
+    PIP_CACHE_DIR=/.cache \
+    WORKERS=1 \
+    THREADS=8
+
+# Update the base OS
+RUN --mount=type=cache,target="/var/cache/apt",sharing=locked \
+    --mount=type=cache,target="/var/lib/apt/lists",sharing=locked \
+    set -eux; \
+    apt-get update; \
+    apt-get upgrade -y; \
+    apt install --no-install-recommends -y  \
+        git; \
+    apt-get autoremove -y
+
+# install base requirements
+COPY requirements-base.txt .
+RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
+    pip install -r requirements-base.txt
+
+# install custom requirements
+COPY requirements.txt .
+RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
+    pip install -r requirements.txt
+
+# install test requirements if needed
+COPY requirements-test.txt .
+# build only when TEST_ENV="true"
+RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
+    if [ "$TEST_ENV" = "true" ]; then \
+      pip install -r requirements-test.txt; \
+    fi
+
+COPY . .
+
+EXPOSE 9090
+
+CMD gunicorn --preload --bind :$PORT --workers $WORKERS --threads $THREADS --timeout 0 _wsgi:app

README.md

@@ -1,10 +1,156 @@
 ---
-title: Ls Huggingface Ner Backend
-emoji: 📉
-colorFrom: blue
-colorTo: indigo
+title: LS Hugging Face NER Backend
+emoji: 🏷️
+colorFrom: purple
+colorTo: pink
 sdk: docker
+app_port: 9090
 pinned: false
+license: apache-2.0
+short_description: HF NER models as a Label Studio ML backend
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LS Hugging Face NER Backend (Hugging Face Spaces)
+
+This Space wraps HumanSignal's [`huggingface_ner` ML backend example](https://github.com/HumanSignal/label-studio-ml-backend/tree/master/label_studio_ml/examples/huggingface_ner) as a Hugging Face Space, so it can be used as a `/predict` backend for a Label Studio instance.
+
+**Patches from the upstream example (minimal):**
+- Added Spaces SDK frontmatter at the top of this README.
+- Removed `docker-compose.yml` (not used on Spaces).
+
+The Dockerfile already binds to `$PORT` and listens on `0.0.0.0`, so it deploys as-is. Default base is `python:3.11-slim` (CPU) — fine for small NER models like `dslim/bert-base-NER`. For larger models, upgrade hardware to a GPU tier.
+
+Connect from Label Studio: set the ML backend URL to `https://davanstrien-ls-huggingface-ner-backend.hf.space`.
+
+---
+
+<!-- Original upstream README below -->
+
+<!--
+---
+title: Hugging Face NER
+type: guide
+tier: all
+order: 25
+hide_menu: true
+hide_frontmatter_title: true
+meta_title: Label Studio tutorial to run Hugging Face NER backend
+meta_description: This tutorial explains how to run a Hugging Face NER backend in Label Studio. 
+categories:
+    - Natural Language Processing
+    - Named Entity Recognition
+    - Hugging Face
+image: "/guide/ml_tutorials/hf-ner.png"
+---
+-->
+
+# Hugging Face NER model with Label Studio
+
+This project uses a custom machine learning backend model for Named Entity Recognition (NER) with Hugging Face's transformers and Label Studio.
+
+The model instantiates `AutoModelForTokenClassification` from Hugging Face's transformers library and fine-tunes it on the NER task.
+
+- If you want to use this model only in inference mode, it serves predictions from the pre-trained model. 
+- If you want to fine-tune the model, you can use the Label Studio interface to provide training data and train the model.
+
+Read more about the compatible models from [Hugging Face's official documentation](https://huggingface.co/docs/transformers/en/tasks/token_classification). 
+
+## Before you begin
+
+Before you begin, you must install the [Label Studio ML backend](https://github.com/HumanSignal/label-studio-ml-backend?tab=readme-ov-file#quickstart). 
+
+This tutorial uses the [`huggingface_ner` example](https://github.com/HumanSignal/label-studio-ml-backend/tree/master/label_studio_ml/examples/huggingface_ner). 
+
+
+## Labeling configuration
+
+This ML backend works with the default NER template from Label Studio. You can find this by selecting Label Studio's pre-built NER template when configuring the labeling interface. It is available under **Natural Language Processing > Named Entity Recognition**:
+
+```xml
+<View>
+  <Labels name="label" toName="text">
+    <Label value="PER" background="red"/>
+    <Label value="ORG" background="darkorange"/>
+    <Label value="LOC" background="orange"/>
+    <Label value="MISC" background="green"/>
+  </Labels>
+
+  <Text name="text" value="$text"/>
+</View>
+```
+
+You can then customize the template to suit your needs (for example, modifying the label names). However, note the model outputs compatibility:
+
+> If you plan to use your model only for the inference, make sure the output label names are compatible with what is listed in XML labeling configuration. If you plan to train the model, you have to provide the baseline pretrained model that can be fine-tuned (i.e. where the last layer can be trained, for example, `distilbert/distilbert-base-uncased`). Otherwise, you may see the error about tensor sizes mismatch during training.
+
+## Running with Docker (recommended)
+
+1. Start the Machine Learning backend on `http://localhost:9090` with the prebuilt image:
+
+```bash
+docker-compose up
+```
+
+2. Validate that backend is running
+
+```bash
+$ curl http://localhost:9090/
+{"status":"UP"}
+```
+
+3. Create a project in Label Studio. Then from the **Model** page in the project settings, [connect the model](https://labelstud.io/guide/ml#Connect-the-model-to-Label-Studio). The default URL is `http://localhost:9090`.
+
+
+## Building from source (advanced)
+
+To build the ML backend from source, you have to clone the repository and build the Docker image:
+
+```bash
+docker-compose build
+```
+
+## Running without Docker (advanced)
+
+To run the ML backend without Docker, you have to clone the repository and install all dependencies using pip:
+
+```bash
+python -m venv ml-backend
+source ml-backend/bin/activate
+pip install -r requirements.txt
+```
+
+Then you can start the ML backend:
+
+```bash
+label-studio-ml start ./huggingface_ner
+```
+
+# Configuration
+
+Parameters can be set in `docker-compose.yml` before running the container.
+
+
+The following common parameters are available:
+- `BASIC_AUTH_USER` - Specify the basic auth user for the model server
+- `BASIC_AUTH_PASS` - Specify the basic auth password for the model server
+- `LOG_LEVEL` - Set the log level for the model server
+- `WORKERS` - Specify the number of workers for the model server
+- `THREADS` - Specify the number of threads for the model server
+- `BASELINE_MODEL_NAME`: The name of the baseline model to use. Default is `dslim/bert-base-NER`.
+- `FINETUNED_MODEL_NAME`: The name of the fine-tuned model. Default is `finetuned_model`.
+- `LABEL_STUDIO_HOST`: The host of the Label Studio instance. Default is 'http://localhost:8080'.
+- `LABEL_STUDIO_API_KEY`: The API key for the Label Studio instance.
+- `START_TRAINING_EACH_N_UPDATES`: The number of updates after which to start training. Default is `10`.
+- `LEARNING_RATE`: The learning rate for the model. Default is `1e-3`.
+- `NUM_TRAIN_EPOCHS`: The number of training epochs. Default is `10`.
+- `WEIGHT_DECAY`: The weight decay for the model. Default is `0.01`.
+- `MODEL_DIR`: The directory where the model is stored. Default is `'./results'`.
+
+> Note: The `LABEL_STUDIO_API_KEY` is required for training the model. This can be found by logging
+  into Label Studio and [going to the **Account & Settings** page](https://labelstud.io/guide/user_account#Access-token). 
+
+# Customization
+
+The ML backend can be customized by adding your own models and logic inside `./huggingface_ner/model.py`.
+
+Modify the `predict()` and `fit()` methods to implement your own logic.

_wsgi.py

@@ -0,0 +1,122 @@
+import os
+import argparse
+import json
+import logging
+import logging.config
+
+logging.config.dictConfig({
+  "version": 1,
+  "disable_existing_loggers": False,
+  "formatters": {
+    "standard": {
+      "format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s"
+    }
+  },
+  "handlers": {
+    "console": {
+      "class": "logging.StreamHandler",
+      "level": os.getenv('LOG_LEVEL'),
+      "stream": "ext://sys.stdout",
+      "formatter": "standard"
+    }
+  },
+  "root": {
+    "level": os.getenv('LOG_LEVEL'),
+    "handlers": [
+      "console"
+    ],
+    "propagate": True
+  }
+})
+
+from label_studio_ml.api import init_app
+from model import HuggingFaceNER
+
+
+_DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json')
+
+
+def get_kwargs_from_config(config_path=_DEFAULT_CONFIG_PATH):
+    if not os.path.exists(config_path):
+        return dict()
+    with open(config_path) as f:
+        config = json.load(f)
+    assert isinstance(config, dict)
+    return config
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Label studio')
+    parser.add_argument(
+        '-p', '--port', dest='port', type=int, default=9090,
+        help='Server port')
+    parser.add_argument(
+        '--host', dest='host', type=str, default='0.0.0.0',
+        help='Server host')
+    parser.add_argument(
+        '--kwargs', '--with', dest='kwargs', metavar='KEY=VAL', nargs='+', type=lambda kv: kv.split('='),
+        help='Additional LabelStudioMLBase model initialization kwargs')
+    parser.add_argument(
+        '-d', '--debug', dest='debug', action='store_true',
+        help='Switch debug mode')
+    parser.add_argument(
+        '--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=None,
+        help='Logging level')
+    parser.add_argument(
+        '--model-dir', dest='model_dir', default=os.path.dirname(__file__),
+        help='Directory where models are stored (relative to the project directory)')
+    parser.add_argument(
+        '--check', dest='check', action='store_true',
+        help='Validate model instance before launching server')
+    parser.add_argument('--basic-auth-user',
+                        default=os.environ.get('ML_SERVER_BASIC_AUTH_USER', None),
+                        help='Basic auth user')
+    
+    parser.add_argument('--basic-auth-pass',
+                        default=os.environ.get('ML_SERVER_BASIC_AUTH_PASS', None),
+                        help='Basic auth pass')    
+    
+    args = parser.parse_args()
+
+    # setup logging level
+    if args.log_level:
+        logging.root.setLevel(args.log_level)
+
+    def isfloat(value):
+        try:
+            float(value)
+            return True
+        except ValueError:
+            return False
+
+    def parse_kwargs():
+        param = dict()
+        for k, v in args.kwargs:
+            if v.isdigit():
+                param[k] = int(v)
+            elif v == 'True' or v == 'true':
+                param[k] = True
+            elif v == 'False' or v == 'false':
+                param[k] = False
+            elif isfloat(v):
+                param[k] = float(v)
+            else:
+                param[k] = v
+        return param
+
+    kwargs = get_kwargs_from_config()
+
+    if args.kwargs:
+        kwargs.update(parse_kwargs())
+
+    if args.check:
+        print('Check "' + HuggingFaceNER.__name__ + '" instance creation..')
+        model = HuggingFaceNER(**kwargs)
+
+    app = init_app(model_class=HuggingFaceNER, basic_auth_user=args.basic_auth_user, basic_auth_pass=args.basic_auth_pass)
+
+    app.run(host=args.host, port=args.port, debug=args.debug)
+
+else:
+    # for uWSGI use
+    app = init_app(model_class=HuggingFaceNER)

model.py

@@ -0,0 +1,260 @@
+import os
+import pathlib
+import re
+import label_studio_sdk
+import logging
+
+from typing import List, Dict, Optional
+from label_studio_ml.model import LabelStudioMLBase
+from label_studio_ml.response import ModelResponse
+from transformers import pipeline, Pipeline
+from itertools import groupby
+from transformers import AutoModelForTokenClassification, TrainingArguments, Trainer, AutoTokenizer
+from transformers import DataCollatorForTokenClassification
+from datasets import Dataset, ClassLabel, Value, Sequence, Features
+from functools import partial
+
+logger = logging.getLogger(__name__)
+_model: Optional[Pipeline] = None
+MODEL_DIR = os.getenv('MODEL_DIR', './results')
+BASELINE_MODEL_NAME = os.getenv('BASELINE_MODEL_NAME', 'dslim/bert-base-NER')
+FINETUNED_MODEL_NAME = os.getenv('FINETUNED_MODEL_NAME', 'finetuned_model')
+
+
+def reload_model():
+    global _model
+    _model = None
+    try:
+        chk_path = str(pathlib.Path(MODEL_DIR) / FINETUNED_MODEL_NAME)
+        logger.info(f"Loading finetuned model from {chk_path}")
+        _model = pipeline("ner", model=chk_path, tokenizer=chk_path)
+    except:
+        # if finetuned model is not available, use the baseline model with the original labels
+        logger.info(f"Loading baseline model {BASELINE_MODEL_NAME}")
+        _model = pipeline("ner", model=BASELINE_MODEL_NAME, tokenizer=BASELINE_MODEL_NAME)
+
+
+reload_model()
+
+
+class HuggingFaceNER(LabelStudioMLBase):
+    """Custom ML Backend model
+    """
+    LABEL_STUDIO_HOST = os.getenv('LABEL_STUDIO_HOST', 'http://localhost:8080')
+    LABEL_STUDIO_API_KEY = os.getenv('LABEL_STUDIO_API_KEY')
+    START_TRAINING_EACH_N_UPDATES = int(os.getenv('START_TRAINING_EACH_N_UPDATES', 10))
+    LEARNING_RATE = float(os.getenv('LEARNING_RATE', 1e-3))
+    NUM_TRAIN_EPOCHS = int(os.getenv('NUM_TRAIN_EPOCHS', 10))
+    WEIGHT_DECAY = float(os.getenv('WEIGHT_DECAY', 0.01))
+
+    def get_labels(self):
+        li = self.label_interface
+        from_name, _, _ = li.get_first_tag_occurence('Labels', 'Text')
+        tag = li.get_tag(from_name)
+        return tag.labels
+    
+    def setup(self):
+        """Configure any paramaters of your model here
+        """
+        self.set("model_version", f'{self.__class__.__name__}-v0.0.1')
+
+    def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs) -> ModelResponse:
+        """ Write your inference logic here
+            :param tasks: [Label Studio tasks in JSON format](https://labelstud.io/guide/task_format.html)
+            :param context: [Label Studio context in JSON format](https://labelstud.io/guide/ml_create#Implement-prediction-logic)
+            :return model_response
+                ModelResponse(predictions=predictions) with
+                predictions: [Predictions array in JSON format](https://labelstud.io/guide/export.html#Label-Studio-JSON-format-of-annotated-tasks)
+        """
+        li = self.label_interface
+        from_name, to_name, value = li.get_first_tag_occurence('Labels', 'Text')
+        texts = [self.preload_task_data(task, task['data'][value]) for task in tasks]
+
+        # run predictions
+        model_predictions = _model(texts)
+
+        predictions = []
+        for prediction in model_predictions:
+            # prediction returned in the format: [{'entity': 'B-ORG', 'score': 0.999, 'index': 1, 'start': 0, 'end': 7, 'word': 'Google'}, ...]
+            # we need to group them by 'B-' and 'I-' prefixes to form entities
+            results = []
+            avg_score = 0
+            for label, group in groupby(prediction, key=lambda x: re.sub(r'^[BI]-', '', x['entity'])):
+                entities = list(group)
+                start = entities[0]['start']
+                end = entities[-1]['end']
+                score = float(sum([entity['score'] for entity in entities]) / len(entities))
+                results.append({
+                    'from_name': from_name,
+                    'to_name': to_name,
+                    'type': 'labels',
+                    'value': {
+                        'start': start,
+                        'end': end,
+                        'labels': [label]
+                    },
+                    'score': score
+                })
+                avg_score += score
+            if results:
+                predictions.append({
+                    'result': results,
+                    'score': avg_score / len(results),
+                    'model_version': self.get('model_version')
+                })
+        
+        return ModelResponse(predictions=predictions, model_version=self.get('model_version'))
+
+    def _get_tasks(self, project_id):
+        # download annotated tasks from Label Studio
+        ls = label_studio_sdk.Client(self.LABEL_STUDIO_HOST, self.LABEL_STUDIO_API_KEY)
+        project = ls.get_project(id=project_id)
+        tasks = project.get_labeled_tasks()
+        return tasks
+
+    def tokenize_and_align_labels(self, examples, tokenizer):
+        """
+        From example https://huggingface.co/docs/transformers/en/tasks/token_classification#preprocess
+        """
+        tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
+
+        labels = []
+        for i, label in enumerate(examples[f"ner_tags"]):
+            word_ids = tokenized_inputs.word_ids(batch_index=i)  # Map tokens to their respective word.
+            previous_word_idx = None
+            label_ids = []
+            for word_idx in word_ids:  # Set the special tokens to -100.
+                if word_idx is None:
+                    label_ids.append(-100)
+                elif word_idx != previous_word_idx:  # Only label the first token of a given word.
+                    label_ids.append(label[word_idx])
+                else:
+                    label_ids.append(-100)
+                previous_word_idx = word_idx
+            labels.append(label_ids)
+
+        tokenized_inputs["labels"] = labels
+        return tokenized_inputs
+    
+    def fit(self, event, data, **kwargs):
+        """Download dataset from Label Studio and prepare data for training in BERT
+        """
+        if event not in ('ANNOTATION_CREATED', 'ANNOTATION_UPDATED', 'START_TRAINING'):
+            logger.info(f"Skip training: event {event} is not supported")
+            return
+
+        # Get project from annotation first if present, otherwise fall back to top-level project field
+        project = data.get('annotation', {}).get('project') or data.get('project')
+        # Handle both possible formats
+        if isinstance(project, dict):
+            project_id = project.get('id')
+        else:
+            project_id = project
+        # If project_id is still None, log and safely exit
+        if project_id is None:
+            logger.error(f"Cannot find project_id in webhook payload: {data}")
+            return
+        
+        tasks = self._get_tasks(project_id)
+
+        if len(tasks) % self.START_TRAINING_EACH_N_UPDATES != 0 and event != 'START_TRAINING':
+            logger.info(f"Skip training: {len(tasks)} tasks are not multiple of {self.START_TRAINING_EACH_N_UPDATES}")
+            return
+
+        # we need to convert Label Studio NER annotations to hugingface NER format in datasets
+        # for example:
+        # {'id': '0',
+        #  'ner_tags': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 7, 8, 8, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0],
+        #  'tokens': ['@paulwalk', 'It', "'s", 'the', 'view', 'from', 'where', 'I', "'m", 'living', 'for', 'two', 'weeks', '.', 'Empire', 'State', 'Building', '=', 'ESB', '.', 'Pretty', 'bad', 'storm', 'here', 'last', 'evening', '.']
+        # }
+        ds_raw = []
+        from_name, to_name, value = self.label_interface.get_first_tag_occurence('Labels', 'Text')
+        tokenizer = AutoTokenizer.from_pretrained(BASELINE_MODEL_NAME)
+
+        no_label = 'O'
+        label_to_id = {no_label: 0}
+        for task in tasks:
+            for annotation in task['annotations']:
+                if not annotation.get('result'):
+                    continue
+                spans = [{'label': r['value']['labels'][0], 'start': r['value']['start'], 'end': r['value']['end']} for r in annotation['result']]
+                spans = sorted(spans, key=lambda x: x['start'])
+                text = self.preload_task_data(task, task['data'][value])
+
+                # insert tokenizer.pad_token to the unlabeled chunks of the text in-between the labeled spans, as well as to the beginning and end of the text
+                last_end = 0
+                all_spans = []
+                for span in spans:
+                    if last_end < span['start']:
+                        all_spans.append({'label': no_label, 'start': last_end, 'end': span['start']})
+                    all_spans.append(span)
+                    last_end = span['end']
+                if last_end < len(text):
+                    all_spans.append({'label': no_label, 'start': last_end, 'end': len(text)})
+
+                # now tokenize chunks separately and add them to the dataset
+                item = {'id': task['id'], 'tokens': [], 'ner_tags': []}
+                for span in all_spans:
+                    tokens = tokenizer.tokenize(text[span['start']:span['end']])
+                    item['tokens'].extend(tokens)
+                    if span['label'] == no_label:
+                        item['ner_tags'].extend([label_to_id[no_label]] * len(tokens))
+                    else:
+                        label = 'B-' + span['label']
+                        if label not in label_to_id:
+                            label_to_id[label] = len(label_to_id)
+                        item['ner_tags'].append(label_to_id[label])
+                        if len(tokens) > 1:
+                            label = 'I-' + span['label']
+                            if label not in label_to_id:
+                                label_to_id[label] = len(label_to_id)
+                            item['ner_tags'].extend([label_to_id[label] for _ in range(1, len(tokens))])
+                ds_raw.append(item)
+
+        logger.debug(f"Dataset: {ds_raw}")
+        # convert to huggingface dataset
+        # Define the features of your dataset
+        features = Features({
+            'id': Value('string'),
+            'tokens': Sequence(Value('string')),
+            'ner_tags': Sequence(ClassLabel(names=list(label_to_id.keys())))
+        })
+        hf_dataset = Dataset.from_list(ds_raw, features=features)
+        tokenized_dataset = hf_dataset.map(partial(self.tokenize_and_align_labels, tokenizer=tokenizer), batched=True)
+
+        logger.debug(f"HF Dataset: {tokenized_dataset}")
+
+        data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)
+        id_to_label = {i: label for label, i in label_to_id.items()}
+        logger.debug(f"Labels: {id_to_label}")
+
+        model = AutoModelForTokenClassification.from_pretrained(
+            BASELINE_MODEL_NAME, num_labels=len(id_to_label),
+            id2label=id_to_label, label2id=label_to_id)
+        logger.debug(f"Model: {model}")
+
+        training_args = TrainingArguments(
+            output_dir=str(pathlib.Path(MODEL_DIR) / FINETUNED_MODEL_NAME),
+            learning_rate=self.LEARNING_RATE,
+            per_device_train_batch_size=8,
+            num_train_epochs=self.NUM_TRAIN_EPOCHS,
+            weight_decay=self.WEIGHT_DECAY,
+            evaluation_strategy="no",
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset,
+            tokenizer=tokenizer,
+            data_collator=data_collator,
+        )
+        trainer.train()
+
+        chk_path = str(pathlib.Path(MODEL_DIR) / FINETUNED_MODEL_NAME)
+        logger.info(f"Model is trained and saved as {chk_path}")
+        trainer.save_model(chk_path)
+
+        # reload model
+        # TODO: this is not thread-safe, should be done with critical section
+        reload_model()

requirements-base.txt

@@ -0,0 +1,2 @@
+gunicorn==23.0.0
+label-studio-ml @ git+https://github.com/HumanSignal/label-studio-ml-backend.git

requirements-test.txt

@@ -0,0 +1,2 @@
+pytest
+pytest-cov

requirements.txt

@@ -0,0 +1,4 @@
+transformers==4.30.2
+datasets==2.18.0
+accelerate==0.28.0
+

test_api.py

@@ -0,0 +1,232 @@
+"""
+This file contains tests for the API of your model. You can run these tests by installing test requirements:
+
+    ```bash
+    pip install -r requirements-test.txt
+    ```
+Then execute `pytest` in the directory of this file.
+
+- Change `NewModel` to the name of the class in your model.py file.
+- Change the `request` and `expected_response` variables to match the input and output of your model.
+"""
+
+import pytest
+import json
+from model import HuggingFaceNER
+import unittest.mock as mock
+
+
+@pytest.fixture
+def client():
+    from _wsgi import init_app
+    app = init_app(model_class=HuggingFaceNER)
+    app.config['TESTING'] = True
+    with app.test_client() as client:
+        yield client
+
+
+def test_predict(client):
+    request = {
+        'tasks': [{
+            'data': {
+                'text': 'President Obama is speaking at 3pm today in New York.'
+            }
+        }],
+        # Your labeling configuration here
+        'label_config': '''
+        <View>
+        <Text name="text" value="$text"/>
+        <Labels name="ner" toName="text">
+            <Label value="Person"/>
+            <Label value="Location"/>
+            <Label value="Time"/>
+        </Labels>
+        </View>
+        '''
+    }
+
+    expected_response = {
+        'results': [{
+            'model_version': 'HuggingFaceNER-v0.0.1',
+            'result': [{
+                'from_name': 'ner',
+                'score': 0.9974774718284607,
+                'to_name': 'text',
+                'type': 'labels',
+                'value': {
+                    'end': 15,
+                    'labels': ['PER'],
+                    'start': 10}},
+                {'from_name': 'ner',
+                 'score': 0.9994751214981079,
+                 'to_name': 'text',
+                 'type': 'labels',
+                 'value': {'end': 52,
+                           'labels': ['LOC'],
+                           'start': 44}}],
+            'score': 0.9984762966632843}]
+    }
+
+    response = client.post('/predict', data=json.dumps(request), content_type='application/json')
+    assert response.status_code == 200
+    response = json.loads(response.data)
+    assert response['results'][0]['model_version'] == expected_response['results'][0]['model_version']
+    assert response['results'][0]['result'][0]['value'] == expected_response['results'][0]['result'][0]['value']
+    assert response['results'][0]['result'][1]['value'] == expected_response['results'][0]['result'][1]['value']
+
+
+# mock response of label_studio_sdk.Project.get_labeled_tasks() and return the list of Label Studio tasks with NER annotations
+def get_labeled_tasks_mock(self, project_id):
+    return [
+        {
+            'id': '0',
+            'data': {'text': 'President Obama is speaking at 3pm today in New York'},
+            'annotations': [
+                {
+                    'result': [
+                        {
+                            'from_name': 'ner',
+                            'to_name': 'text',
+                            'type': 'labels',
+                            'value': {
+                                'start': 10,
+                                'end': 15,
+                                'labels': ['Person']
+                            }
+                        },
+                        {
+                            'from_name': 'ner',
+                            'to_name': 'text',
+                            'type': 'labels',
+                            'value': {
+                                'start': 44,
+                                'end': 52,
+                                'labels': ['Location']
+                            }
+                        },
+                        {
+                            'from_name': 'ner',
+                            'to_name': 'text',
+                            'type': 'labels',
+                            'value': {
+                                'start': 31,
+                                'end': 40,
+                                'labels': ['Time']
+                            }
+                        }
+                    ]
+                }
+            ]
+        }
+    ]
+
+
+# mock NewModel.START_TRAINING_EACH_N_UPDATES to 1 to trigger training in the test
+@pytest.fixture
+def mock_start_training():
+    with mock.patch.object(HuggingFaceNER, 'START_TRAINING_EACH_N_UPDATES', new=1):
+        yield
+
+
+@pytest.fixture
+def mock_get_labeled_tasks():
+    with mock.patch.object(HuggingFaceNER, '_get_tasks', new=get_labeled_tasks_mock):
+        yield
+
+
+@pytest.fixture
+def mock_baseline_model_name_for_train():
+    with mock.patch('model.BASELINE_MODEL_NAME', new='distilbert/distilbert-base-uncased'):
+        yield
+
+
+def test_fit(client, mock_get_labeled_tasks, mock_start_training, mock_baseline_model_name_for_train):
+    request = {
+        'action': 'ANNOTATION_CREATED',
+        'project': {
+            'id': 12345,
+            'label_config': '''
+        <View>
+        <Text name="text" value="$text"/>
+        <Labels name="ner" toName="text">
+            <Label value="Person"/>
+            <Label value="Location"/>
+            <Label value="Time"/>
+        </Labels>
+        </View>
+        '''
+        },
+        'annotation': {
+            'project': 12345
+        }
+    }
+
+    response = client.post('/webhook', data=json.dumps(request), content_type='application/json')
+    assert response.status_code == 201
+
+    # assert new model is created in ./results/finetuned_model directory
+    import os
+    from model import MODEL_DIR
+    results_dir = os.path.join(MODEL_DIR, 'finetuned_model')
+    assert os.path.exists(os.path.join(results_dir, 'pytorch_model.bin'))
+
+    # now let's test whether the model is trained by running predict
+    request = {
+        'tasks': [{
+            'data': {
+                'text': 'President Obama is speaking at 3pm today in New York.'
+            }
+        }],
+        # Your labeling configuration here
+        'label_config': '''
+        <View>
+        <Text name="text" value="$text"/>
+        <Labels name="ner" toName="text">
+            <Label value="Person"/>
+            <Label value="Location"/>
+            <Label value="Time"/>
+        </Labels>
+        </View>
+        '''
+    }
+
+    response = client.post('/predict', data=json.dumps(request), content_type='application/json')
+    assert response.status_code == 200
+
+    # TODO: we also need to check the prediction results to make sure the model is trained correctly
+    # but the training needs to be deterministic to make the test stable
+    # assert response is as expected
+
+    # remove './results/finetuned_model' directory after testing
+    import shutil
+    shutil.rmtree(results_dir)
+
+def test_fit_missing_annotation(monkeypatch):
+    # Initialize the model
+    model = HuggingFaceNER()
+
+    # Mock label_interface to avoid AttributeError
+    model.label_interface = mock.MagicMock()
+    # Mock get_first_tag_occurence to return fake values
+    model.label_interface.get_first_tag_occurence.return_value = ('Labels', 'Text', 'text_field_name')
+
+    # Mock data payload with annotation missing, only project present
+    payload = {
+        "action": "ANNOTATION_UPDATED",
+        "project": {"id": 123, "name": "Test Project"}
+    }
+
+    # Monkeypatch _get_tasks to return one fake task
+    monkeypatch.setattr(model, "_get_tasks", lambda project_id: [
+        {
+            "id": "1",
+            "data": {"text_field_name": "Hello world"},
+            "annotations": []
+        }
+    ])
+
+    # Call fit()
+    try:
+        model.fit(event="ANNOTATION_UPDATED", data=payload)
+    except Exception as e:
+        pytest.fail(f"fit() raised an exception when annotation is missing: {e}")