new task: pointer network for joint ner and re

[ArneBinder]

code is mostly taken from https://github.com/ArneBinder/pie-document-level/pull/26

Note: This adds the requirement transformers = "^4.32.0", so that we can import BartPreTrainedModel (breaking).

Requires:

• pytorch-ie>=0.29.5 (add TaskModule.configure_model_metric(stage) pytorch-ie#392)
• new task: text-2-text #24 (for the PIE model)
• re-arrange and add taskmodule metrics #25 (for metrics)
• add BartModelWithDecoderPositionIds base model #26 (for decoder_position_id_pattern, i.e. replaced positional encodings (RPE))

TODO:

• taskmodule: add code and get tests passing
• taskmodule: understand by refactoring (x + 3h + 8h)
• model: add code and get tests passing (prediction: 8h, true: 5h)
• model: understand by refactoring (estimation: 16h, true: 3h+?)
  • cleanup ids and get old setup running (true: 3h)
  • implement simplified version (estimation: 8h, true: 42h):
    • use HF generate for pointer network #20 (estimation: 8h, true: 28h)
    • no RPE (estimation: 2h, true: 1h)
    • mask train labels in simple model (estimation: 1h, true: 1h)
    • understand and fix weight decay issues (true: 6h)
    • find a good way what to log when (finally, add metric_intervals parameter) (true: 2h)
    • show that training is aligned (true: 4h)
• use pytorch-ie = ">=0.29.5,<0.30.0" when released (we require add TaskModule.configure_model_metric(stage) pytorch-ie#392) - [x] constrained training (true: 4h)
• taskmodule
  • merge PointerNetworkSpanAndRelationEncoderDecoder back again into taskmodule and improve structure (true: 10h)
  • there are a lot of missed annotations during tokenization... EDIT: should be fixed (we did not set return_overflowing_tokens=True)
  • check if cmp_src_rel sorts the relations in a useful manner (note: we encode relations as tail-head-label). Edit: we sort first by head, then by tail. We could try the other way, which should be more in sync with the encoding mode, but the current setup seems to work.
  • understand sanitize_sequence
  • tests
    • increase coverage
      • add tests for common and pointer_network modules (true: 2h)
      • add tests for sanitize_sequence()
    • simplify / disentangle existing ones
  • rename (we tackle end2end_re with the current taskmodule variant) . Edit: fixed in 3411b29
  • generalize metrics
    • add / improve tests
  • use PrefixConstrainedLogitsProcessor to constrain the generation (true: 3h + 5h)
• model:
  • tests
    • increase coverage
    • add test_pointer_head.py (PointerHead)
    • improve test_bart_as_pointer_network (BartAsPointerNetwork)
    • remove test_simple_generative_pointer.py
    • remove test_simple_generative_pointer_predict.py
    • check if fixture data is unused and delete it if that is the case
    • use mock models? Edit: use sshleifer/tiny-mbart use sshleifer/bart-tiny-random
  • improve weight decay parametrization (currently we use a single parameter layernorm_decay to set "encoder only layernorm parameters"). Edit: fixed in b8da7e3
  • fix metric logging warning. Edit: fixed in 0070649
  • move original model into separate branch. Edit: done, see https://github.com/ArneBinder/pie-modules/tree/original_pointer_network
  • parametrize if to use predictions or model output to calculate metrics (the latter should be much faster, but overestimating). Implemented in 6b01b1b
  • implement RPE. Edit: finalized in d5facf9
  • generalize PIE model to GenerativeModel. Edit: see new task: text-2-text #24
  • outsource BartModelWithDecoderPositionIds, see add BartModelWithDecoderPositionIds base model #26
  • move prepare_decoder_position_ids to BartAsPointerNetwork Edit: no, this is fine
  • use special_tokens_mask in PointerHead Edit: no, just mask eos, bos, and pad positions in offset scores

Follow-up:

• adapt other base models, see transformers.AutoModelForSeq2SeqLM for candidates
  • e.g. LongT5ForConditionalGeneration (pszemraj/long-t5-tglobal-base-sci-simplify)
• maybe: use BucketSampler, see Allowed custom BatchSamplers when instantiated in *_dataloader hook Lightning-AI/pytorch-lightning#13640 (comment)

[codecov-commenter]

Codecov Report

Attention: 43 lines in your changes are missing coverage. Please review.

Comparison is base (7198a39) 95.64% compared to head (b79aa68) 95.56%.

Files	Patch %	Lines
...ules/taskmodules/pointer_network_for_end2end_re.py	93.58%	25 Missing ⚠️
...ules/models/base_models/bart_as_pointer_network.py	93.42%	10 Missing ⚠️
src/pie_modules/models/components/pointer_head.py	97.65%	3 Missing ⚠️
...ules/pointer_network/annotation_encoder_decoder.py	97.98%	3 Missing ⚠️
src/pie_modules/taskmodules/common/interfaces.py	88.88%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #10      +/-   ##
==========================================
- Coverage   95.64%   95.56%   -0.09%     
==========================================
  Files          34       40       +6     
  Lines        2459     3315     +856     
==========================================
+ Hits         2352     3168     +816     
- Misses        107      147      +40     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ArneBinder]

It looks like the training is now aligned. When we replace the weights of the original model with the ones from the simple model, we get the following:

# check the numbers
assert {layer_name: len(anns) for layer_name, anns in fp.items()} == {
    "labeled_spans": 101,
    "binary_relations": 63,
}
assert {layer_name: len(anns) for layer_name, anns in fn.items()} == {
    "labeled_spans": 93,
    "binary_relations": 73,
}
assert {layer_name: len(anns) for layer_name, anns in tp.items()} == {
    "labeled_spans": 41,
    "binary_relations": 9,
}

which is quite close to the result from the not-modified original model:

# check the numbers
assert {layer_name: len(anns) for layer_name, anns in fp.items()} == {
    "labeled_spans": 126,
    "binary_relations": 72,
}
assert {layer_name: len(anns) for layer_name, anns in fn.items()} == {
    "labeled_spans": 88,
    "binary_relations": 73,
}
assert {layer_name: len(anns) for layer_name, anns in tp.items()} == {
    "labeled_spans": 46,
    "binary_relations": 9,
}

The result for the simple model we took the weights from is much worse:

# check the numbers
assert {layer_name: len(anns) for layer_name, anns in fp.items()} == {
    "labeled_spans": 108,
    "binary_relations": 65,
}
assert {layer_name: len(anns) for layer_name, anns in fn.items()} == {
    "labeled_spans": 106,
    "binary_relations": 78,
}
assert {layer_name: len(anns) for layer_name, anns in tp.items()} == {
    "labeled_spans": 29,
    "binary_relations": 4,
}