Fix a crazy bug and add README

Some tokens
after the punctuation has been removed become joint into a single token,
for instance @.@ becomes @@,
which we then cannot find in the original string.
So switch from using find,
to using the end indices
and keeping track of punctuation we removed from the original string.

README.md

@@ -10,26 +10,116 @@ pinned: true
 app_port: 8000
 ---
 
+# Comma fixer
+This repository contains a web service for fixing comma placement within a given text, for instance:
+
+`"A sentence however, not quite good correct and sound."` -> `"A sentence, however, not quite good, correct and sound."`
+
+It provides a webpage for testing the functionality, a REST API,
+and Jupyter notebooks for evaluating and training comma fixing models.
+
+A web demo is hosted in the [huggingface spaces](https://huggingface.co/spaces/klasocki/comma-fixer).
+
+## Development setup
+
+Deploying the service for local development can be done by running `docker-compose up` in the root directory.
+Note that you might have to
 `sudo service docker start`
+first.
+
+The application should then be available at http://localhost:8000.
+For the API, see the `openapi.yaml` file.
+Docker-compose mounts a volume and listens to changes in the source code, so the application will be reloaded and 
+reflect them.
+
+We use multi-stage builds to reduce the image size, ensure flexibility in requirements and that tests are run before 
+each deployment.
+However, while it does reduce the size by nearly 3GB, the resulting image still contains deep learning libraries and 
+pre-downloaded models, and will take around 7GB of disk space.
+
+Alternatively, you can setup a python environment by hand. It is recommended to use a virtualenv. Inside one, run
+```bash
+pip install -e .[test]
+```
+the `[test]` option makes sure to install test dependencies.
+
+If you intend to perform training and evaluation of deep learning models, install also using the `[training]` option. 
+
+### Running tests
+To run the tests, execute
+```bash
+docker build -t comma-fixer --target test .
+``` 
+Or `python -m pytest tests/ ` if you already have a local python environment.
 
-`docker log [id]` for logs from the container.
 
-`docker build -t comma-fixer --target test .` for tests
+### Deploying to huggingface spaces
+In order to deploy the application, one needs to be added as a collaborator to the space and have set up a
+corresponding git remote.
+The application is then continuously deployed on each push.
+```bash
+git remote add hub https://huggingface.co/spaces/klasocki/comma-fixer
+git push hub
+```
 
-`git push hub` to deploy to huggingface hub, after adding a remote
+## Evaluation
 
-Multi-stage build brings down the size from 9GB+ to around 7GB.
-Less not possible most likely, due to the size of torch and models.
+In order to evaluate, run `jupyter notebook notebooks/` or copy the notebooks to a web hosting service with GPUs, 
+such as Google Colab or Kaggle
+and clone this repository there.
 
-Reported token classification F1 scores on commas for different languages, on a political speeches' dataset:
+We use the [oliverguhr/fullstop-punctuation-multilang-large](https://huggingface.co/oliverguhr/fullstop-punctuation-multilang-large)
+model as the baseline.
+It is a RoBERTa large model fine-tuned for the task of punctuation restoration on a dataset of political speeches
+in English, German, French and Italian.
+That is, it takes a sentence without any punctuation as input, and predicts the missing punctuation in token 
+classification fashion, thanks to which the original token structure stays unchanged.
+We use a subset of its capabilities focusing solely on commas, and leaving other punctuation intact.
+
+
+
+The authors report the following token classification F1 scores on commas for different languages on the original 
+dataset:
 
 | English | German | French | Italian |
 |---------|--------|--------|---------|
 | 0.819   | 0.945  | 0.831  | 0.798   |
 
-Evaluation of the baseline model on the wikitext-103-raw-v1 validation dataset:
+The results of our evaluation of the baseline model out of domain on the English wikitext-103-raw-v1 validation 
+dataset are as follows:
 
 | precision | recall | F1   | support |
 |-----------|--------|------|---------|
 | 0.79      | 0.71   | 0.75 | 10079   |
 
+We treat each comma as one token instance, as opposed to the original paper, which NER-tags the whole multiple-token 
+preceding words as comma class tokens.
+In our approach, for each comma from the prediction text obtained from the model:
+*  If it should be there according to ground truth, it counts as a true positive.
+ * If it should not be there, it counts as a false positive.
+ * If a comma from ground truth is not predicted, it counts as a false negative.
+
+## Training
+While fine-tuning an encoder BERT-like pre-trained model for NER seems like the best approach to the problem, 
+since it preserves the sentence structure and only focuses on commas,
+with limited GPU resources, we doubt we could beat the baseline model with a similar approach.
+We could fine-tune the baseline on our data, focusing on commas, and see if it brings any improvement.
+
+However, we thought that trying out pre-trained text-to-text or decoder-only LLMs for this task using PEFT could be 
+interesting, and wanted to check if we have enough resources for low-rank adaptation or prefix-tuning.
+
+We adapt the code from [this tutorial](https://www.youtube.com/watch?v=iYr1xZn26R8) in order to fine-tune a 
+[bloom LLM](https://huggingface.co/bigscience/bloom-560m) to our task using 
+[LoRa](https://arxiv.org/pdf/2106.09685.pdf).
+However, even with the smallest model from the family, we struggled with CUDA memory errors using the free Google 
+colab GPU quotas, and could only train with a batch size of two.
+After a short training, it seems the loss keeps fluctuating and the model is only able to learn to repeat the 
+original phrase back. 
+
+If time permits, we plan to experiment with seq2seq pre-trained models, increasing gradient accumulation steps, and the 
+percentage of 
+data with commas.
+The latter could help since wikitext contains highly diverse data, with many rows being empty strings, 
+headers, or short paragraphs.
+
+

commafixer/src/baseline.py

@@ -1,52 +1,131 @@
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline, NerPipeline
+import re
 
 
 class BaselineCommaFixer:
+    """
+    A wrapper class for the oliverguhr/fullstop-punctuation-multilang-large baseline punctuation restoration model.
+    It adapts the model to perform comma fixing instead of full punctuation restoration, that is, removes the
+    punctuation, runs the model, and then uses its outputs so that only commas are changed.
+    """
+
     def __init__(self, device=-1):
         self._ner = _create_baseline_pipeline(device=device)
 
     def fix_commas(self, s: str) -> str:
+        """
+        The main method for fixing commas using the baseline model.
+        In the future we should think about batching the calls to it, for now it processes requests string by string.
+        :param s: A string with commas to fix, without length restrictions.
+        Example: comma_fixer.fix_commas("One two thre, and four!")
+        :return: A string with commas fixed, example: "One, two, thre and four!"
+        """
+        s_no_punctuation, punctuation_indices = _remove_punctuation(s)
         return _fix_commas_based_on_pipeline_output(
-            self._ner(_remove_punctuation(s)),
-            s
+            self._ner(s_no_punctuation),
+            s,
+            punctuation_indices
         )
 
 
 def _create_baseline_pipeline(model_name="oliverguhr/fullstop-punctuation-multilang-large", device=-1) -> NerPipeline:
+    """
+    Creates the huggingface pipeline object.
+    Can also be used for pre-downloading the model and the tokenizer.
+    :param model_name: Name of the baseline model on the huggingface hub.
+    :param device: Device to use when running the pipeline, defaults to -1 for CPU, a higher number indicates the id
+    of GPU to use.
+    :return: A token classification pipeline.
+    """
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForTokenClassification.from_pretrained(model_name)
     return pipeline('ner', model=model, tokenizer=tokenizer, device=device)
 
 
-def _remove_punctuation(s: str) -> str:
-    to_remove = ".,?-:"
-    for char in to_remove:
-        s = s.replace(char, '')
-    return s
-
-
-def _fix_commas_based_on_pipeline_output(pipeline_json: list[dict], original_s: str) -> str:
+def _remove_punctuation(s: str) -> tuple[str, list[int]]:
+    """
+    Removes the punctuation (".,?-:") from the input text, since the baseline model has been trained on data without
+    punctuation. It also keeps track of the indices where we remove it, so that we can restore the original later.
+    Commas are the exception, since we remove them, but restore with the model.
+    Hence we do not keep track of removed comma indices.
+    :param s: For instance, "A short-string: with punctuation, removed.
+    :return: A tuple of a string, for instance:
+    "A shortstring with punctuation removed"; and a list of indices where punctuation has been removed, in ascending
+    order
+    """
+    to_remove_regex = r"[\.\?\-:]"
+    # We're not counting commas, since we will remove them later anyway. Only counting removals that will be restored
+    # in the final resulting string.
+    punctuation_indices = [m.start() for m in re.finditer(to_remove_regex, s)]
+    s = re.sub(to_remove_regex, '', s)
+    s = s.replace(',', '')
+    return s, punctuation_indices
+
+
+def _fix_commas_based_on_pipeline_output(pipeline_json: list[dict], original_s: str, punctuation_indices: list[int]) -> \
+        str:
+    """
+    This function takes the comma fixing token classification pipeline output, and converts it to string based on the
+    original
+    string and punctuation indices, so that the string contains all the original characters, except commas, intact.
+    :param pipeline_json: Token classification pipeline output.
+    Contains five fields.
+    'entity' is the punctuation that should follow this token.
+    'word' is the token text together with preceding space if any.
+    'end' is the end index in the original string (with punctuation removed in our case!!)
+    Example: [{'entity': ':',
+  'score': 0.90034866,
+  'index': 1,
+  'word': '▁Exam',
+  'start': 0,
+  'end': 4},
+ {'entity': ':',
+  'score': 0.9157294,
+  'index': 2,
+  'word': 'ple',
+  'start': 4,
+  'end': 7}]
+    :param original_s: The original string, before removing punctuation.
+    :param punctuation_indices: The indices of the removed punctuation except commas, so that we can correctly keep
+    track of the current offset in the original string.
+    :return: A string with commas fixed, and other the original punctuation from the input string restored.
+    """
     result = original_s.replace(',', '')  # We will fix the commas, but keep everything else intact
-    current_offset = 0
+
+    commas_inserted_or_punctuation_removed = 0
+    removed_punctuation_index = 0
 
     for i in range(1, len(pipeline_json)):
-        current_offset = _find_current_token(current_offset, i, pipeline_json, result)
+        current_offset = pipeline_json[i - 1]['end'] + commas_inserted_or_punctuation_removed
+
+        commas_inserted_or_punctuation_removed, current_offset, removed_punctuation_index = (
+            _update_offset_by_the_removed_punctuation(
+                commas_inserted_or_punctuation_removed, current_offset, punctuation_indices, removed_punctuation_index
+            )
+        )
+
         if _should_insert_comma(i, pipeline_json):
             result = result[:current_offset] + ',' + result[current_offset:]
-            current_offset += 1
+            commas_inserted_or_punctuation_removed += 1
     return result
 
 
-def _should_insert_comma(i, pipeline_json, new_word_indicator='▁') -> bool:
-    # Only insert commas for the final token of a word
-    return pipeline_json[i - 1]['entity'] == ',' and pipeline_json[i]['word'].startswith(new_word_indicator)
+def _update_offset_by_the_removed_punctuation(
+        commas_inserted_and_punctuation_removed, current_offset, punctuation_indices, removed_punctuation_index
+):
+    # increase the counters for every punctuation removed from the original string before the curent offset
+    while (removed_punctuation_index < len(punctuation_indices) and
+           punctuation_indices[removed_punctuation_index] < current_offset):
 
+        commas_inserted_and_punctuation_removed += 1
+        removed_punctuation_index += 1
+        current_offset += 1
+    return commas_inserted_and_punctuation_removed, current_offset, removed_punctuation_index
 
-def _find_current_token(current_offset, i, pipeline_json, result, new_word_indicator='▁') -> int:
-    current_word = pipeline_json[i - 1]['word'].replace(new_word_indicator, '')
-    # Find the current word in the result string, starting looking at current offset
-    current_offset = result.find(current_word, current_offset) + len(current_word)
-    return current_offset
+
+def _should_insert_comma(i, pipeline_json, new_word_indicator='▁') -> bool:
+    # Only insert commas for the final token of a word, that is, if next word starts with a space.
+    return pipeline_json[i - 1]['entity'] == ',' and pipeline_json[i]['word'].startswith(new_word_indicator)
 
 
 if __name__ == "__main__":

notebooks/evaluation.ipynb

@@ -283,15 +283,6 @@
     }
    ]
   },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "We have 2 commas predicted correctly, and 3 missed, so we are expecting 100% precision and 40% recall."
-   ],
-   "metadata": {
-    "id": "NzVo05UcoPlb"
-   }
-  },
   {
    "cell_type": "code",
    "source": [
@@ -346,6 +337,15 @@
     }
    ]
   },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "We have 2 commas predicted correctly, and 3 missed, so we are expecting 100% precision and 40% recall."
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
   {
    "cell_type": "code",
    "source": [

openapi.yaml

@@ -32,11 +32,10 @@ paths:
                   s:
                     type: string
                     example: 'This is a sentence with wrong commas, at least some.'
-                    description: A text with commas fixed, or unchanged if not necessary. 
-                      TODO WARNING - the text will have spaces normalized and trimmed at the start and end.
-                      TODO some other punctuation may be changed as well
+                    description: A text with commas fixed, or unchanged if not necessary. Everything other that 
+                      commas will stay as it was originally.
 
         400:
-          description: Input text query parameter missing.
+          description: A required field missing from the POST request body JSON.
 
   

setup.py

@@ -21,7 +21,7 @@ setup(
     extras_require={
         'training': [
             'datasets==2.14.4',
-            'seqeval'
+            'seqeval',
             'notebook'
         ],
         'test': [

tests/test_baseline.py

@@ -30,7 +30,15 @@ def test_fix_commas_leaves_correct_strings_unchanged(baseline_fixer, test_input)
         ['Even newlines\ntabs\tand others get preserved.',
          'Even newlines,\ntabs\tand others get preserved.'],
         ['I had no Creativity left, therefore, I come here, and write useless examples, for this test.',
-         'I had no Creativity left therefore, I come here and write useless examples for this test.']]
+         'I had no Creativity left therefore, I come here and write useless examples for this test.'],
+        [' This is a sentence. With, a lot of, useless punctuation!!??. O.o However we have to insert commas O-O, '
+         'nonetheless or we will fail this test.',
+         ' This is a sentence. With a lot of useless punctuation!!??. O.o However, we have to insert commas O-O '
+         'nonetheless, or we will fail this test.'],
+        [" The ship 's secondary armament consisted of fourteen 45 @-@ calibre 6 @-@ inch ( 152 mm ) quick @-@ firing ( QF ) guns mounted in casemates . Lighter guns consisted of eight 47 @-@ millimetre ( 1 @.@ 9 in ) three @-@ pounder Hotchkiss guns and four 47 @-@ millimetre 2 @.@ 5 @-@ pounder Hotchkiss guns . The ship was also equipped with four submerged 18 @-@ inch torpedo tubes two on each broadside .",
+         " The ship 's secondary armament consisted of fourteen 45 @-@ calibre 6 @-@ inch ( 152 mm ) quick @-@ firing ( QF ) guns mounted in casemates . Lighter guns consisted of eight 47 @-@ millimetre ( 1 @.@ 9 in ), three @-@ pounder Hotchkiss guns and four 47 @-@ millimetre 2 @.@ 5 @-@ pounder Hotchkiss guns . The ship was also equipped with four submerged 18 @-@ inch torpedo tubes, two on each broadside ."]
+
+    ]
 )
 def test_fix_commas_fixes_incorrect_commas(baseline_fixer, test_input, expected):
     result = baseline_fixer.fix_commas(s=test_input)
@@ -39,10 +47,11 @@ def test_fix_commas_fixes_incorrect_commas(baseline_fixer, test_input, expected)
 
 @pytest.mark.parametrize(
     "test_input, expected",
-    [['', ''],
-     ['Hello world...', 'Hello world'],
+    [['', ('', [])],
+     [' world...', (' world', [6, 7, 8])],
+     [',,,', ('', [])],
      ['This: test - string should not, have any commas inside it...?',
-      'This test  string should not have any commas inside it']]
+      ('This test  string should not have any commas inside it', [4, 11, 57, 58, 59, 60])]]
 )
 def test__remove_punctuation(test_input, expected):
     assert _remove_punctuation(test_input) == expected