backend added

.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,128 @@
+<!--
 ---
-title: Rlhf Docker
-emoji: 🦀
-colorFrom: blue
-colorTo: indigo
-sdk: docker
-pinned: false
+title: Classify text with a BERT model
+type: guide
+tier: all
+order: 35
+hide_menu: true
+hide_frontmatter_title: true
+meta_title: BERT-based text classification
+meta_description: Tutorial on how to use BERT-based text classification with your Label Studio project
+categories:
+    - Natural Language Processing
+    - Text Classification
+    - BERT
+    - Hugging Face
+image: "/tutorials/bert.png"
 ---
+-->
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# BERT-based text classification
+
+The NewModel is a BERT-based text classification model that is designed to work with Label Studio. This model uses the Hugging Face Transformers library to fine-tune a BERT model for text classification. The model is trained on the labeled data from Label Studio and then used to make predictions on new data.  With this model connected to Label Studio, you can: 
+
+- Train a BERT model on your labeled data directly from Label Studio.
+- Use any model for [AutoModelForSequenceClassification](https://huggingface.co/transformers/v3.0.2/model_doc/auto.html#automodelforsequenceclassification) from the Hugging Face model hub.
+- Fine-tune the model on your specific task and use it to make predictions on new data.
+- Automatically download the labeled tasks from Label Studio and prepare the data for training.
+- Customize the training parameters such as learning rate, number of epochs, and weight decay.
+
+## 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 [`bert_classifier` example](https://github.com/HumanSignal/label-studio-ml-backend/tree/master/label_studio_ml/examples/bert_classifier). 
+
+
+## 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`.
+
+> Warning! Note the current limitation of the ML backend: models are loaded dynamically from huggingface.co. You may need the `HF_TOKEN` env variable provided in your environment. Consequently, this may result in a slow response time for the first prediction request. If you are experiencing timeouts on Label Studio side (i.e., no predictions are visible when opening the task), check the logs of the ML backend for any errors, and refresh the page in a few minutes.
+
+## 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 ./dir_with_your_model
+```
+
+
+## Labeling configuration
+
+In project `Settings > Labeling Interface > Browse Templates > Natural Language Processing > Text Classification`, you can find the default labeling configuration for text classification in Label Studio. This configuration includes a single `<Choices>` output and a single `<Text>` input. 
+Feel free to modify the set of labels in the `<Choices>` tag to match your specific task, for example:
+
+```xml
+<View>
+  <Text name="text" value="$text" />
+  <Choices name="label" toName="text" choice="single" showInLine="true">
+    <Choice value="label one" />
+    <Choice value="label two" />
+      <Choice value="label three" />
+  </Choices>
+</View>
+```
+
+
+## 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 for training. Default is `bert-base-multilingual-cased`.
+
+## Training
+
+The following parameters are available for training:
+
+- `LABEL_STUDIO_HOST` (required): The URL of the Label Studio instance. Default is `http://localhost:8080`.
+- `LABEL_STUDIO_API_KEY` (required): The [API key](https://labelstud.io/guide/user_account#Access-token) for the Label Studio instance.
+- `START_TRAINING_EACH_N_UPDATES`: The number of labeled tasks to download from Label Studio before starting training. Default is 10.
+- `LEARNING_RATE`: The learning rate for the model training. Default is 2e-5.
+- `NUM_TRAIN_EPOCHS`: The number of epochs for model training. Default is 3.
+- `WEIGHT_DECAY`: The weight decay for the model training. Default is 0.01.
+- `FINETUNED_MODEL_NAME`: The name of the fine-tuned model. Default is `finetuned_model`. Checkpoints will be saved under this name.
+
+> Note: The `LABEL_STUDIO_API_KEY` is required for training the model. You can find the API key in Label Studio under 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 the `./bert_classifier` directory.

_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 BertClassifier
+
+
+_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 "' + BertClassifier.__name__ + '" instance creation..')
+        model = BertClassifier(**kwargs)
+
+    app = init_app(model_class=BertClassifier, 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=BertClassifier)

docker-compose.yml

@@ -0,0 +1,44 @@
+version: "3.8"
+
+services:
+  bert_classifier:
+    container_name: bert_classifier
+    image: heartexlabs/label-studio-ml-backend:bertclass-master
+    build:
+      context: .
+      args:
+        TEST_ENV: ${TEST_ENV}
+    environment:
+      # If you are using this model for training, you have to connect it to Label Studio
+      - LABEL_STUDIO_HOST=http://localhost:8080
+      - LABEL_STUDIO_API_KEY=your-api-key
+      # Use any model for [AutoModelForSequenceClassification](https://huggingface.co/transformers/v3.0.2/model_doc/auto.html#automodelforsequenceclassification)
+      - BASELINE_MODEL_NAME=bert-base-multilingual-cased
+      # - BASELINE_MODEL_NAME=google/electra-small-discriminator
+      # The model directory for the fine-tuned checkpoints (relative to $MODEL_DIR)
+      - FINETUNED_MODEL_NAME=finetuned_model
+      # The number of labeled tasks to download from Label Studio before starting training
+      - START_TRAINING_EACH_N_UPDATES=10
+      # Learning rate
+      - LEARNING_RATE=2e-5
+      # Number of epochs
+      - NUM_TRAIN_EPOCHS=3
+      # Weight decay
+      - WEIGHT_DECAY=0.01
+      # specify these parameters if you want to use basic auth for the model server
+      - BASIC_AUTH_USER=
+      - BASIC_AUTH_PASS=
+      # set the log level for the model server
+      - LOG_LEVEL=DEBUG
+      # any other parameters that you want to pass to the model server
+      - ANY=PARAMETER
+      # specify the number of workers and threads for the model server
+      - WORKERS=1
+      - THREADS=8
+      # specify the model directory (likely you don't need to change this)
+      - MODEL_DIR=/data/models
+    ports:
+      - "9090:9090"
+    volumes:
+      - "./data/server:/data"
+      - "./data/.cache:/root/.cache"

model.py

@@ -0,0 +1,214 @@
+import os
+import torch
+import logging
+import pathlib
+import label_studio_sdk
+
+from typing import List, Dict, Optional
+from label_studio_ml.model import LabelStudioMLBase
+from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments
+from transformers import pipeline
+from label_studio_sdk.label_interface.objects import PredictionValue
+from label_studio_ml.response import ModelResponse
+from datasets import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print('There are %d GPU(s) available.' % torch.cuda.device_count())
+    print('We will use the GPU:', torch.cuda.get_device_name(0))
+else:
+    print('No GPU available, using the CPU instead.')
+    device = torch.device("cpu")
+
+
+class BertClassifier(LabelStudioMLBase):
+    """
+    BERT-based text classification model for Label Studio
+
+    This model uses the Hugging Face Transformers library to fine-tune a BERT model for text classification.
+    Use any model for [AutoModelForSequenceClassification](https://huggingface.co/transformers/v3.0.2/model_doc/auto.html#automodelforsequenceclassification)
+    The model is trained on the labeled data from Label Studio and then used to make predictions on new data.
+
+    Parameters:
+    -----------
+    LABEL_STUDIO_HOST : str
+        The URL of the Label Studio instance
+    LABEL_STUDIO_API_KEY : str
+        The API key for the Label Studio instance
+    START_TRAINING_EACH_N_UPDATES : int
+        The number of labeled tasks to download from Label Studio before starting training
+    LEARNING_RATE : float
+        The learning rate for the model training
+    NUM_TRAIN_EPOCHS : int
+        The number of epochs for model training
+    WEIGHT_DECAY : float
+        The weight decay for the model training
+    baseline_model_name : str
+        The name of the baseline model to use for training
+    MODEL_DIR : str
+        The directory to save the trained model
+    finetuned_model_name : str
+        The name of the finetuned 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', 2e-5))
+    NUM_TRAIN_EPOCHS = int(os.getenv('NUM_TRAIN_EPOCHS', 3))
+    WEIGHT_DECAY = float(os.getenv('WEIGHT_DECAY', 0.01))
+    baseline_model_name = os.getenv('BASELINE_MODEL_NAME', 'bert-base-multilingual-cased')
+    MODEL_DIR = os.getenv('MODEL_DIR', './results')
+    finetuned_model_name = os.getenv('FINETUNED_MODEL_NAME', 'finetuned-model')
+    _model = None
+
+    def get_labels(self):
+        li = self.label_interface
+        from_name, _, _ = li.get_first_tag_occurence('Choices', 'Text')
+        tag = li.get_tag(from_name)
+        return tag.labels
+
+    def setup(self):
+        self.set("model_version", f'{self.__class__.__name__}-v0.0.1')
+
+    def _lazy_init(self):
+        if not self._model:
+            try:
+                chk_path = str(pathlib.Path(self.MODEL_DIR) / self.finetuned_model_name)
+                self._model = pipeline("text-classification", model=chk_path, tokenizer=chk_path)
+            except:
+                # if finetuned model is not available, use the baseline model, with the labels from the label_interface
+                self._model = pipeline(
+                    "text-classification",
+                    model=self.baseline_model_name,
+                    tokenizer=self.baseline_model_name)
+
+                labels = self.get_labels()
+                self._model.model.config.id2label = {i: label for i, label in enumerate(labels)}
+                self._model.model.config.label2id = {label: i for i, label in enumerate(labels)}
+
+    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 predictions: [Predictions array in JSON format](https://labelstud.io/guide/export.html#Label-Studio-JSON-format-of-annotated-tasks)
+        """
+
+        # TODO: this may result in single-time timeout for large models - consider adjusting the timeout on Label Studio side
+        self._lazy_init()
+
+        li = self.label_interface
+        from_name, to_name, value = li.get_first_tag_occurence('Choices', 'Text')
+        texts = [self.preload_task_data(task, task['data'][value]) for task in tasks]
+
+        model_predictions = self._model(texts)
+        predictions = []
+        for prediction in model_predictions:
+            logger.debug(f"Prediction: {prediction}")
+            region = li.get_tag(from_name).label(prediction['label'])
+            pv = PredictionValue(
+                score=prediction['score'],
+                result=[region],
+                model_version=self.get('model_version')
+            )
+            predictions.append(pv)
+
+        return ModelResponse(predictions=predictions)
+
+    def fit(self, event, data, **additional_params):
+        """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
+        project_id = data['annotation']['project']
+
+        # dowload 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()
+
+        logger.info(f"Downloaded {len(tasks)} labeled tasks from Label Studio")
+        logger.debug(f"Tasks: {tasks}")
+        if len(tasks) % self.START_TRAINING_EACH_N_UPDATES != 0 and event != 'START_TRAINING':
+            # skip training if the number of tasks is not divisible by START_TRAINING_EACH_N_UPDATES
+            logger.info(f"Skip training: the number of tasks is not divisible by {self.START_TRAINING_EACH_N_UPDATES}")
+            return
+
+        from_name, to_name, value = self.label_interface.get_first_tag_occurence('Choices', 'Text')
+
+        ds_raw = {
+            'id': [],
+            'text': [],
+            'label': []
+        }
+        for task in tasks:
+            for annotation in task['annotations']:
+                if 'result' in annotation:
+                    for result in annotation['result']:
+                        if 'choices' in result['value']:
+                            ds_raw['id'].append(task['id'])
+                            text = self.preload_task_data(task, task['data'][value])
+                            ds_raw['text'].append(text)
+                            ds_raw['label'].append(result['value']['choices'])
+
+        hf_dataset = Dataset.from_dict(ds_raw)
+        logger.debug(f"Dataset: {hf_dataset}")
+
+        labels = self.get_labels()
+        label_to_id = {label: i for i, label in enumerate(labels)}
+        id_to_label = {i: label for i, label in enumerate(labels)}
+        logger.debug(f"Labels: {labels}")
+
+        # Preprocess the dataset
+        tokenizer = AutoTokenizer.from_pretrained(self.baseline_model_name)
+
+        def preprocess_function(examples):
+            return tokenizer(examples["text"], truncation=True, padding=True)
+
+        tokenized_datasets = hf_dataset.map(preprocess_function, batched=True)
+        logger.debug(f"Tokenized dataset: {tokenized_datasets}")
+
+        # Convert labels to ids
+        def label_to_id_function(examples):
+            examples["label"] = [label_to_id[label] for label in examples["label"]]
+            return examples
+
+        tokenized_datasets = tokenized_datasets.map(label_to_id_function)
+
+        # Load model with custom config
+        logger.info(f"Start training the model {self.finetuned_model_name}")
+        config = AutoConfig.from_pretrained(self.baseline_model_name, num_labels=len(labels))
+        logger.debug(f"Config: {config}")
+        model = AutoModelForSequenceClassification.from_pretrained(self.baseline_model_name, config=config)
+        model.config.id2label = id_to_label
+        model.config.label2id = label_to_id
+        logger.debug(f"Model: {model}")
+
+        # Define training arguments
+        training_args = TrainingArguments(
+            output_dir=str(pathlib.Path(self.MODEL_DIR) / 'training_output'),
+            learning_rate=2e-5,
+            evaluation_strategy="no",
+            num_train_epochs=3,
+            weight_decay=0.01,
+            log_level='info'
+        )
+        logger.debug(f"Training arguments: {training_args}")
+
+        # Initialize Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_datasets,
+            tokenizer=tokenizer,
+        )
+        logger.debug(f"Trainer: {trainer}")
+
+        # Train the model
+        trainer.train()
+
+        chk_path = str(pathlib.Path(self.MODEL_DIR) / self.finetuned_model_name)
+        logger.info(f"Model is trained and saved as {chk_path}")
+        trainer.save_model(chk_path)

requirements-base.txt

@@ -0,0 +1,2 @@
+gunicorn==22.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,3 @@
+transformers==4.38.0
+datasets==2.18.0
+accelerate==0.28.0

test_api.py

@@ -0,0 +1,62 @@
+"""
+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 BertClassifier
+
+
+@pytest.fixture
+def client():
+    from _wsgi import init_app
+    app = init_app(model_class=BertClassifier)
+    app.config['TESTING'] = True
+    with app.test_client() as client:
+        yield client
+
+
+def test_predict(client):
+    request = {
+        'tasks': [{
+            'data': {
+                'text': 'Today is a great day to play football.'
+            }
+        }],
+        # Your labeling configuration here
+        'label_config':
+            '<View>'
+            '<Text name="text" value="$text" />'
+            '<Choices name="topic" toName="text" choice="single">'
+            '<Choice value="sports" />'
+            '<Choice value="politics" />'
+            '<Choice value="technology" />'
+            '</Choices>'
+            '</View>'
+    }
+
+    expected_response_results = [{
+        'result': [{
+            'from_name': 'topic',
+            'to_name': 'text',
+            'type': 'choices',
+            'value': {'choices': ['sports']}
+        }]
+    }]
+
+    response = client.post('/predict', data=json.dumps(request), content_type='application/json')
+    assert response.status_code == 200
+    response = json.loads(response.data)
+    assert len(expected_response_results) == len(response['results'])
+
+
+# TODO
+# Implement test_fit()