Sync from 0bb0e6d in FinetuningBert on GitHub (#7)

- Sync from 66b3fde in FinetuningBert on GitHub (026c61a877cdec918ece3e1411b8b785afb653b1)
- Fix trust_remote_code loading + multi-label dependency head indexing (9d0eccc9a093e4984899a51245dcc1005c291b4f)

BertForJointParsing.py

@@ -7,12 +7,22 @@ from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
 from transformers.models.bert.modeling_bert import BertOnlyMLMHead
 from transformers.utils import ModelOutput
-from .BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits
-from .BertForPrefixMarking import BertPrefixMarkingHead, parse_logits as prefix_parse_logits, encode_sentences_for_bert_for_prefix_marking, get_prefixes_from_str
-from .BertForMorphTagging import BertMorphTaggingHead, MorphLogitsOutput, MorphLabels, parse_logits as morph_parse_logits 
+# transformers' get_imports() strips the try body but keeps the except body, so the absolute
+# imports have to sit in the try or trust_remote_code demands them as pip packages
+try:
+    from BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits
+    from BertForPrefixMarking import BertPrefixMarkingHead, parse_logits as prefix_parse_logits, encode_sentences_for_bert_for_prefix_marking, get_prefixes_from_str
+    from BertForMorphTagging import BertMorphTaggingHead, MorphLogitsOutput, MorphLabels, parse_logits as morph_parse_logits
+except ImportError:
+    from .BertForSyntaxParsing import BertSyntaxParsingHead, SyntaxLabels, SyntaxLogitsOutput, parse_logits as syntax_parse_logits
+    from .BertForPrefixMarking import BertPrefixMarkingHead, parse_logits as prefix_parse_logits, encode_sentences_for_bert_for_prefix_marking, get_prefixes_from_str
+    from .BertForMorphTagging import BertMorphTaggingHead, MorphLogitsOutput, MorphLabels, parse_logits as morph_parse_logits
     
 import warnings
 
+INDEX_TO_LABELS = ['syntax', 'ner', 'prefix', 'lex', 'morph']
+LABELS_TO_INDEX = {k:i for i,k in enumerate(INDEX_TO_LABELS)}
+
 @dataclass
 class JointParsingOutput(ModelOutput):
     loss: Optional[torch.FloatTensor] = None
@@ -45,7 +55,7 @@ class BertForJointParsing(BertPreTrainedModel):
         "cls.predictions.decoder.bias": "cls.predictions.bias",
     }
 
-    def __init__(self, config, do_syntax=None, do_ner=None, do_prefix=None, do_lex=None, do_morph=None, syntax_head_size=64):
+    def __init__(self, config, do_syntax=None, do_ner=None, do_prefix=None, do_lex=None, do_morph=None, syntax_head_size=128):
         super().__init__(config)
 
         self.bert = BertModel(config, add_pooling_layer=False)
@@ -95,7 +105,7 @@ class BertForJointParsing(BertPreTrainedModel):
         position_ids: Optional[torch.Tensor] = None,
         prefix_class_id_options: Optional[torch.Tensor] = None,
         labels: Optional[Union[SyntaxLabels, MorphLabels, torch.Tensor]] = None,
-        labels_type: Optional[Literal['syntax', 'ner', 'prefix', 'lex', 'morph']] = None,
+        labels_type: Optional[torch.Tensor] = None,
         head_mask: Optional[torch.Tensor] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
         output_attentions: Optional[bool] = None,
@@ -109,7 +119,8 @@ class BertForJointParsing(BertPreTrainedModel):
         if labels is not None and labels_type is None:
             raise ValueError("Cannot specify labels without labels_type")
         
-        if labels_type == 'seg' and prefix_class_id_options is None:
+        labels_type_str = INDEX_TO_LABELS[labels_type.detach().cpu().item()] if labels_type is not None else None
+        if labels_type_str == 'seg' and prefix_class_id_options is None:
             raise ValueError('Cannot calculate prefix logits without prefix_class_id_options')
         
         if compute_syntax_mst is not None and self.syntax is None:
@@ -144,13 +155,13 @@ class BertForJointParsing(BertPreTrainedModel):
         morph_logits = None
 
         # Calculate the syntax
-        if self.syntax is not None and (labels is None or labels_type == 'syntax'):
+        if self.syntax is not None and (labels is None or labels_type_str == 'syntax'):
             # apply the syntax head
             loss, syntax_logits = self.syntax(hidden_states, extended_attention_mask, labels, compute_syntax_mst)
             logits = syntax_logits
 
         # Calculate the NER
-        if self.ner is not None and (labels is None or labels_type == 'ner'):
+        if self.ner is not None and (labels is None or labels_type_str == 'ner'):
             ner_logits = self.ner(hidden_states)
             logits = ner_logits
             if labels is not None:
@@ -158,19 +169,19 @@ class BertForJointParsing(BertPreTrainedModel):
                 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
 
         # Calculate the segmentation
-        if self.prefix is not None and (labels is None or labels_type == 'prefix'):
+        if self.prefix is not None and (labels is None or labels_type_str == 'prefix'):
             loss, prefix_logits = self.prefix(hidden_states, prefix_class_id_options, labels)
             logits = prefix_logits
         
         # Calculate the lexeme
-        if self.lex is not None and (labels is None or labels_type == 'lex'):
+        if self.lex is not None and (labels is None or labels_type_str == 'lex'):
             lex_logits = self.lex(hidden_states)
             logits = lex_logits
             if labels is not None:
                 loss_fct = nn.CrossEntropyLoss()  # -100 index = padding token
                 loss = loss_fct(lex_logits.view(-1, self.config.vocab_size), labels.view(-1))
 
-        if self.morph is not None and (labels is None or labels_type == 'morph'):
+        if self.morph is not None and (labels is None or labels_type_str == 'morph'):
             loss, morph_logits = self.morph(hidden_states, labels)
             logits = morph_logits
 
@@ -215,7 +226,7 @@ class BertForJointParsing(BertPreTrainedModel):
         final_output = [dict(text=sentence, tokens=combine_token_wordpieces(ids, offsets, tokenizer)) for sentence, ids, offsets in zip(sentences, input_ids, offset_mapping)]
         # Syntax logits: each sentence gets a dict(tree: List[dict(word,dep_head,dep_head_idx,dep_func)], root_idx: int)
         if output.syntax_logits is not None:
-            for sent_idx,parsed in enumerate(syntax_parse_logits(input_ids, sentences, tokenizer, output.syntax_logits)):
+            for sent_idx,parsed in enumerate(syntax_parse_logits(input_ids, sentences, tokenizer, output.syntax_logits, self.syntax.func_label_idx)):
                 merge_token_list(final_output[sent_idx]['tokens'], parsed['tree'], 'syntax')
                 final_output[sent_idx]['root_idx'] = parsed['root_idx']
                 
@@ -404,7 +415,7 @@ def convert_output_to_ud(output_sentences, model_cfg, style: Literal['htb', 'iah
                 for pre in get_prefixes_from_str(word['seg'][0], model_cfg.prefix_cfg, greedy=True):
                     # pos - just take the first valid pos that appears in the predicted prefixes list. 
                     pos = next((pos for pos in ud_prefixes_to_pos[pre] if pos in word['morph']['prefixes']), ud_prefixes_to_pos[pre][0])
-                    dep, func = ud_get_prefix_dep(pre, word, word_idx)
+                    dep, func = ud_get_prefix_dep(pre, word, word_idx, sentence['tokens'])
                     intermediate_output.append(dict(word=pre, lex=pre, pos=pos, dep=dep, func=func, feats='_'))
                     
                 # if there was an implicit heh, add it in dependent on the method
@@ -464,7 +475,6 @@ def convert_output_to_ud(output_sentences, model_cfg, style: Literal['htb', 'iah
                     # if the function is obj, then add a את pointing to the next word
                     elif func == 'obj' and s_lex != 'את':
                         intermediate_output.append(dict(word='_את_', lex='את', pos='ADP', dep=len(intermediate_output) + 2, func='case', feats='_', absolute_dep=True))
-
                 # add the main suffix in
                 intermediate_output.append(dict(word=suf, lex='הוא', pos='PRON', dep=dep, func=func, feats='|'.join(f'{k}={v}' for k,v in word['morph']['suffix_feats'].items())))
                 if entry_to_assign_suf_dep:
@@ -512,13 +522,32 @@ def normalize_dep_rel(dep, style: Literal['htb', 'iahlt']):
     return dep
     
     
-def ud_get_prefix_dep(pre, word, word_idx):
+def ud_get_prefix_dep(pre, word, word_idx, all_toks):
     does_follow_main = False
+    is_head_prev = False 
+    does_follow_prev = False
     
+    # Special rabbinic daled case as an adposition, like psukei *de*-zimra
+    if pre == "ד" and "ADP" in word['morph']['prefixes']:
+        func = "case:gen"
+    # Special rabbinic quf case as an aspect-tense marker, like *qa*-amar
+    elif pre == "ק":
+        func = "aux"
     # shin goes to the main word for verbs, otherwise follows the word
-    if pre.endswith('ש'): 
-        does_follow_main = word['morph']['pos'] != 'VERB' and word['syntax']['dep_head_idx'] > word_idx
-        func = 'mark'
+    elif pre.endswith('ש') or pre.endswith('ד'): 
+        # short circuit for fixed subjugators (e.g. לפני ש, אף אל פי ש)
+        func = None
+        if word_idx > 0 and all_toks[word_idx - 1]['syntax']['dep_func'] == 'mark':
+            is_head_prev = True
+            func = 'fixed'
+        elif word_idx > 1 and all_toks[word_idx - 1]['syntax']['dep_func'] == 'fixed':
+            grandparent_index = all_toks[word_idx - 1]['syntax']['dep_head_idx']
+            if all_toks[grandparent_index]['syntax']['dep_func'] == 'mark':
+                does_follow_prev = True
+                func = 'fixed'
+        if func is None:
+            does_follow_main = word['morph']['pos'] != 'VERB' and word['syntax']['dep_head_idx'] > word_idx
+            func = 'mark'
     # vuv goes to the main word if the function is in the list, otherwise follows
     elif pre == 'ו': 
         does_follow_main = word['syntax']['dep_func'] not in ["conj", "acl:recl", "parataxis", "root", "acl", "amod", "list", "appos", "dep", "flatccomp"]
@@ -534,4 +563,8 @@ def ud_get_prefix_dep(pre, word, word_idx):
         if pre == 'ה':
             func = 'det' if 'DET' in word['morph']['prefixes'] else 'mark'
 
+    if is_head_prev:
+        return word_idx - 1, func
+    elif does_follow_prev:
+        return all_toks[word_idx - 1]['syntax']['dep_head_idx'], func
     return (word['syntax']['dep_head_idx'] if does_follow_main else word_idx), func

BertForMorphTagging.py

@@ -16,6 +16,12 @@ ALL_FEATURES = [
     ('Person', ['none', '1', '2', '3', '1,2,3']),
     ('Tense', ['none', 'Past', 'Fut', 'Pres', 'Imp'])
 ]
+ALL_FEATURES_MAP = {
+    'Gender': {'Fem, Masc': 'Fem,Masc'},
+    'Number': {},
+    'Person': {},
+    'Tense': { 'Present': 'Pres' }
+}
 
 @dataclass
 class MorphLogitsOutput(ModelOutput):
@@ -55,19 +61,22 @@ class BertMorphTaggingHead(nn.Module):
         super().__init__()
         self.config = config
 
-        self.num_prefix_classes = len(ALL_PREFIX_POS)
-        self.num_pos_classes = len(ALL_POS)
-        self.num_suffix_classes = len(ALL_SUFFIX_POS)
-        self.num_features_classes = list(map(len, map(itemgetter(1), ALL_FEATURES)))
+        if not hasattr(config, 'classes_config'):
+            config.classes_config = dict(ALL_PREFIX_POS=ALL_PREFIX_POS, ALL_POS=ALL_POS, ALL_SUFFIX_POS=ALL_SUFFIX_POS, ALL_FEATURES=list(map(itemgetter(1), ALL_FEATURES)))
+        
+        self.num_prefix_classes = len(config.classes_config['ALL_PREFIX_POS'])
+        self.num_pos_classes = len(config.classes_config['ALL_POS'])
+        self.num_suffix_classes = len(config.classes_config['ALL_SUFFIX_POS'])
+        self.num_features_classes = list(map(len, config.classes_config['ALL_FEATURES']))
         # we need a classifier for prefix cls and POS cls
         # the prefix will use BCEWithLogits for multiple labels cls
         self.prefix_cls = nn.Linear(config.hidden_size, self.num_prefix_classes)
         # and pos + feats will use good old cross entropy for single label
         self.pos_cls = nn.Linear(config.hidden_size, self.num_pos_classes)
-        self.features_cls = nn.ModuleList([nn.Linear(config.hidden_size, len(features)) for _, features in ALL_FEATURES])
+        self.features_cls = nn.ModuleList([nn.Linear(config.hidden_size, len(features)) for features in config.classes_config['ALL_FEATURES']])
         # and suffix + feats will also be cross entropy
         self.suffix_cls = nn.Linear(config.hidden_size, self.num_suffix_classes)
-        self.suffix_features_cls = nn.ModuleList([nn.Linear(config.hidden_size, len(features)) for _, features in ALL_FEATURES])
+        self.suffix_features_cls = nn.ModuleList([nn.Linear(config.hidden_size, len(features)) for features in config.classes_config['ALL_FEATURES']])
 
     def forward(
             self, 
@@ -84,7 +93,7 @@ class BertMorphTaggingHead(nn.Module):
         if labels is not None:
             # step 1: prefix labels loss
             loss_fct = nn.BCEWithLogitsLoss(weight=(labels.prefix_labels != -100).float())
-            loss = loss_fct(prefix_logits, labels.prefix_labels)
+            loss = loss_fct(prefix_logits, labels.prefix_labels) * 25#is the *25 needed? No idea, but worked for Rabbinic
             # step 2: pos labels loss
             loss_fct = nn.CrossEntropyLoss()
             loss += loss_fct(pos_logits.view(-1, self.num_pos_classes), labels.pos_labels.view(-1))
@@ -178,7 +187,6 @@ def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: Be
     special_toks = tokenizer.all_special_tokens
     special_toks.remove(tokenizer.unk_token)
     special_toks.remove(tokenizer.mask_token)
-    
     ret = []
     for sent_idx,sentence in enumerate(sentences):
         input_id_strs = tokenizer.convert_ids_to_tokens(input_ids[sent_idx])

BertForPrefixMarking.py

@@ -1,9 +1,15 @@
 from transformers.utils import ModelOutput
 import torch
 from torch import nn
-from typing import Dict, List, Tuple, Optional
+from typing import Dict, List, Tuple, Optional, Union
 from dataclasses import dataclass
-from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
+from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast, AutoConfig, AutoModel
+from transformers.models.auto.auto_factory import _BaseAutoModelClass
+try:
+    from transformers.modeling_utils import no_init_weights
+except ImportError:
+    from transformers.initialization import no_init_weights
+import inspect, os
 
 # define the classes, and the possible prefixes for each class
 POSSIBLE_PREFIX_CLASSES =  [ ['לכש', 'כש', 'מש', 'בש', 'לש'], ['מ'], ['ש'], ['ה'], ['ו'], ['כ'], ['ל'], ['ב'] ]
@@ -104,14 +110,22 @@ class BertPrefixMarkingHead(nn.Module):
         
         return (loss, logits)
         
-
-
-class BertForPrefixMarking(BertPreTrainedModel):
-
-    def __init__(self, config):
+def can_func_take_parameter(fn, param_name):
+    signature = inspect.signature(fn)
+    # Exclude 'self' from parameters
+    parameters = [p.name for p in signature.parameters.values() if p.name != 'self']
+    return 'kwargs' in parameters or param_name in parameters
+
+class BaseForPrefixMarking(BertPreTrainedModel):
+    base_model_prefix = ""
+    def __init__(self, config, bert_cls=BertModel):
         super().__init__(config)
+        setattr(config, "hidden_dropout_prob", getattr(config, "hidden_dropout_prob", 0.1))
+        setattr(config, "initializer_range", getattr(config, "classifier_init_range", getattr(config, 'decoder_init_range', 0.02)))
+
+        self.bert = bert_cls(config, **({} if not can_func_take_parameter(bert_cls.__init__, 'add_pooling_layer') else {'add_pooling_layer': False}))
+        self.send_token_type_ids = can_func_take_parameter(self.bert.forward, 'token_type_ids')
 
-        self.bert = BertModel(config, add_pooling_layer=False)
         self.dropout = nn.Dropout(config.hidden_dropout_prob)
         self.prefix = BertPrefixMarkingHead(config)
 
@@ -138,16 +152,17 @@ class BertForPrefixMarking(BertPreTrainedModel):
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        kwargs = dict(token_type_ids=token_type_ids, head_mask=head_mask) if self.send_token_type_ids else {}
+
         bert_outputs = self.bert(
             input_ids,
             attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
             position_ids=position_ids,
-            head_mask=head_mask,
             inputs_embeds=inputs_embeds,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            **kwargs
         )
 
         hidden_states = bert_outputs[0]
@@ -174,6 +189,20 @@ class BertForPrefixMarking(BertPreTrainedModel):
         logits = self.forward(**inputs, return_dict=True).logits
         return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits, self.config.prefix_cfg)
 
+
+class AutoForPrefixMarking(_BaseAutoModelClass):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *args, **kwargs):
+        auto_cfg = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        base_cls = BaseForPrefixMarking
+        with no_init_weights():
+            bert_cls = AutoModel.from_config(auto_cfg, *args, **{k: v for k, v in kwargs.items() if k != 'config'}).__class__
+            if 'Prefix' in bert_cls.__name__:
+                base_cls = bert_cls
+                
+        return base_cls.from_pretrained(pretrained_model_name_or_path, *args, **kwargs, bert_cls=bert_cls, key_mapping={"^model": "bert"})
+    
+
 def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: torch.FloatTensor, config: PrefixConfig):
     # extract the predictions by argmaxing the final dimension (batch x sequence x prefixes x prediction)
     logit_preds = torch.argmax(logits, axis=3).tolist()
@@ -182,6 +211,7 @@ def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: Be
 
     for sent_idx,sent_ids in enumerate(input_ids):
         tokens = tokenizer.convert_ids_to_tokens(sent_ids)
+        
         ret.append([])
         for tok_idx,token in enumerate(tokens):
             # If we've reached the pad token, then we are at the end
@@ -193,6 +223,9 @@ def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: Be
             while next_tok_idx < len(tokens) and tokens[next_tok_idx].startswith('##'):
                 token += tokens[next_tok_idx][2:]
                 next_tok_idx += 1
+            
+            if hasattr(tokenizer, 'splinter') and tokenizer.splinter:
+                token = tokenizer.splinter.unsplinter_word(token)
 
             prefix_len = get_predicted_prefix_len_from_logits(token, logit_preds[sent_idx][tok_idx], config)
         

BertForSyntaxParsing.py

@@ -1,12 +1,19 @@
-import math
+import math, os, inspect
 from transformers.utils import ModelOutput
 import torch
 from torch import nn
 from typing import Dict, List, Tuple, Optional, Union
 from dataclasses import dataclass
-from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast
+from transformers import BertPreTrainedModel, BertModel, BertTokenizerFast, AutoModel, AutoConfig
+from transformers.models.auto.auto_factory import _BaseAutoModelClass
+try:
+    from transformers.modeling_utils import no_init_weights
+except ImportError:
+    from transformers.initialization import no_init_weights
+
 
 ALL_FUNCTION_LABELS = ["nsubj", "nsubj:cop", "punct", "mark", "mark:q", "case", "case:gen", "case:acc", "fixed", "obl", "det", "amod", "acl:relcl", "nmod", "cc", "conj", "root", "compound:smixut", "cop", "compound:affix", "advmod", "nummod", "appos", "nsubj:pass", "nmod:poss", "xcomp", "obj", "aux", "parataxis", "advcl", "ccomp", "csubj", "acl", "obl:tmod", "csubj:pass", "dep", "dislocated", "nmod:tmod", "nmod:npmod", "flat", "obl:npmod", "goeswith", "reparandum", "orphan", "list", "discourse", "iobj", "vocative", "expl", "flat:name"]
+ALL_POS = ['DET', 'NOUN', 'VERB', 'CCONJ', 'ADP', 'PRON', 'PUNCT', 'ADJ', 'ADV', 'SCONJ', 'NUM', 'PROPN', 'AUX', 'X', 'INTJ', 'SYM']
 
 @dataclass
 class SyntaxLogitsOutput(ModelOutput):
@@ -20,7 +27,7 @@ class SyntaxLogitsOutput(ModelOutput):
 @dataclass
 class SyntaxTaggingOutput(ModelOutput):
     loss: Optional[torch.FloatTensor] = None
-    logits: Optional[SyntaxLogitsOutput] = None
+    logits: Optional[Union[torch.FloatTensor, SyntaxLogitsOutput]] = None
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
     attentions: Optional[Tuple[torch.FloatTensor]] = None
 
@@ -28,13 +35,124 @@ class SyntaxTaggingOutput(ModelOutput):
 class SyntaxLabels(ModelOutput):
     dependency_labels: Optional[torch.LongTensor] = None
     function_labels: Optional[torch.LongTensor] = None
+    pos_labels: Optional[torch.LongTensor] = None
 
     def detach(self):
-        return SyntaxLabels(self.dependency_labels.detach(), self.function_labels.detach())
+        return SyntaxLabels(self.dependency_labels.detach(), self.function_labels.detach(), self.pos_labels.detach() if self.pos_labels is not None else None)
+    
+    def to(self, device, non_blocking=False):
+        return SyntaxLabels(self.dependency_labels.to(device, non_blocking=non_blocking), self.function_labels.to(device, non_blocking=non_blocking), self.pos_labels.to(device, non_blocking=non_blocking) if self.pos_labels is not None else None)
+
+class BertSyntaxPartialInfoHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        # we want an embedding table of size 64 for each function label & unk + a single parameter embedding for unknown head
+        FUNCTION_CLASS_EMBED_SIZE = 64
+        TRANSFORM_SIZE = config.hidden_size * 2 + FUNCTION_CLASS_EMBED_SIZE
+        self.function_class_embed = nn.Embedding(len(ALL_FUNCTION_LABELS) + 1, FUNCTION_CLASS_EMBED_SIZE)
+        self.unk_function_class = len(ALL_FUNCTION_LABELS) # 0-based
+        self.head_unk_embed = nn.Embedding(1, config.hidden_size)
+
+        if False:
+            TRANSFORM_SIZE += FUNCTION_CLASS_EMBED_SIZE
+            self.pos_class_embed = nn.Embedding(len(ALL_POS) + 1, FUNCTION_CLASS_EMBED_SIZE)
+            self.unk_pos_class = len(ALL_POS) # 0-based
+
+        # Linear layer to transform the hidden states + activation
+        self.transform = nn.Linear(TRANSFORM_SIZE, config.hidden_size)
+        self.activation = nn.Tanh()
+
+        # Auxiliary classifier to predict the input function labels from the transformed hidden states
+        if False:
+            self.aux_function_classifier = nn.Linear(config.hidden_size, len(ALL_FUNCTION_LABELS))
+            self.aux_loss_weight = getattr(config, 'partial_info_aux_loss_weight', 0.6)
+        
+        # Storage for auxiliary losses from each layer (will be accumulated during forward passes)
+        # Note: This won't work correctly in distributed training - we check and skip if distributed
+        self.aux_losses = []
+        self._partial_labels = None
+        if torch.distributed.is_initialized():
+            raise NotImplementedError("Partial info head not supported in distributed training")
     
-    def to(self, device):
-        return SyntaxLabels(self.dependency_labels.to(device), self.function_labels.to(device))
+    def clear_aux_losses(self):
+        self.aux_losses = []
+
+    def set_partial_labels(self, partial_labels: SyntaxLabels):
+        self._partial_labels = partial_labels
+
+    def get_total_aux_loss(self):
+        if not self.aux_losses:
+            return None
+        total_loss = sum(self.aux_losses)
+        return total_loss * self.aux_loss_weight
+
+    def forward(
+        self, 
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        
+        is_tuple, tuple_values = False, (None,)
+        if isinstance(hidden_states, tuple):
+            is_tuple = True
+            hidden_states, tuple_values = hidden_states[0], hidden_states[1:]
+
+        # lookup the function embeddings - turn the -1s into the unknown function label
+        function_embeddings = self.function_class_embed(replace_tensor_value(self._partial_labels.function_labels.clamp_min(-1), -1, self.unk_function_class))
+        # lookup the dependency embeddings - first just lookup the labels in the hidden_states - for now -1 gets clamped to 0, we don't care
+        # after that, replace the -1s, with the actual value
+        dependency_embedding = torch.gather(hidden_states, 1, self._partial_labels.dependency_labels.unsqueeze(-1).expand(-1, -1, self.config.hidden_size).clamp_min(0))
+        dependency_embedding = torch.where((self._partial_labels.dependency_labels == -1).unsqueeze(-1), self.head_unk_embed.weight[0], dependency_embedding)
+
+        # cat them all into a single embedding
+        intermediate_states = torch.cat([hidden_states, function_embeddings, dependency_embedding], dim=-1)
+
+        if False:
+            pos_embeddings = self.pos_class_embed(replace_tensor_value(self._partial_labels.pos_labels.clamp_min(-1), -1, self.unk_pos_class))
+            intermediate_states = torch.cat([intermediate_states, pos_embeddings], dim=-1)
+
+        # run through transform and activation
+        transformed = self.activation(self.transform(intermediate_states))
+
+        # Auxiliary classifier: predict the input function labels from the transformed hidden states
+        if self.training and self._partial_labels is not None and False:
+            aux_logits = self.aux_function_classifier(transformed)
+            # Compute auxiliary loss - only on positions where we have valid function labels (not -1)
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
+            aux_loss = loss_fct(aux_logits.view(-1, len(ALL_FUNCTION_LABELS)), self._partial_labels.function_labels.clamp_min(-1).view(-1))
+            self.aux_losses.append(aux_loss)
+        
+        if is_tuple:
+            return (transformed, *tuple_values)
+        return transformed
+
+class BertSyntaxValidClassifierHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+
+        self.transform = nn.Linear(config.hidden_size, config.hidden_size)
+        self.act = nn.Tanh()
+        self.cls = nn.Linear(config.hidden_size, 2) # valid / invalid
+
+    def forward(
+        self, 
+        hidden_states: torch.Tensor, 
+        extended_attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        compute_mst: bool = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        # transform the hidden states
+        transformed_states = self.act(self.transform(hidden_states[:, 0, :])) # batch x dim
+        logits = self.cls(transformed_states) # batch x 2
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, 2), labels.view(-1))
 
+        return (loss, logits)
+        
 class BertSyntaxParsingHead(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -42,11 +160,15 @@ class BertSyntaxParsingHead(nn.Module):
 
         # the attention query & key values
         self.head_size = config.syntax_head_size# int(config.hidden_size / config.num_attention_heads * 2)
-        self.query = nn.Linear(config.hidden_size, self.head_size)
-        self.key = nn.Linear(config.hidden_size, self.head_size)
+        self.label_count = getattr(config, 'syntax_attn_label_count', 1)
+        self.query = nn.Linear(config.hidden_size, self.head_size * self.label_count)
+        self.key = nn.Linear(config.hidden_size, self.head_size * self.label_count)
         # the function classifier gets two encoding values and predicts the labels
+        self.func_label_idx = getattr(config, 'syntax_func_label_idx', 0)
         self.num_function_classes = len(ALL_FUNCTION_LABELS)
-        self.cls = nn.Linear(config.hidden_size * 2, self.num_function_classes)
+        if self.func_label_idx > -1:
+            self.cls = nn.Linear(config.hidden_size * 2, self.num_function_classes)
+        else: self.cls = None
 
     def forward(
             self, 
@@ -54,36 +176,44 @@ class BertSyntaxParsingHead(nn.Module):
             extended_attention_mask: Optional[torch.Tensor],
             labels: Optional[SyntaxLabels] = None,
             compute_mst: bool = False) -> Tuple[torch.Tensor, SyntaxLogitsOutput]:
+        if compute_mst:
+            assert self.label_count == 1, "Cannot compute MST with multiple attention labels - please set syntax_attn_label_count to 1"
         
-        # Take the dot product between "query" and "key" to get the raw attention scores.
-        query_layer = self.query(hidden_states)
-        key_layer = self.key(hidden_states)
-        attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)) / math.sqrt(self.head_size)
+        device = hidden_states.device
 
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        hidden_shape = (*hidden_states.shape[:-1], -1, self.head_size) # batch x seq x label_count x head_size
+        query_layer = self.query(hidden_states).view(*hidden_shape).transpose(1, 2)
+        key_layer = self.key(hidden_states).view(*hidden_shape).transpose(1, 2)
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) / math.sqrt(self.head_size) # batch x label_count x seq x seq
+        
         # add in the attention mask
         if extended_attention_mask is not None:
-            if extended_attention_mask.ndim == 4:
-                extended_attention_mask = extended_attention_mask.squeeze(1)
-            attention_scores += extended_attention_mask# batch x seq x seq
+            attention_scores += extended_attention_mask# batch x label_count x seq x seq
 
         # At this point take the hidden_state of the word and of the dependency word, and predict the function
         # If labels are provided, use the labels.
         if self.training and labels is not None:
             # Note that the labels can have -100, so just set those to zero with a max
-            dep_indices = labels.dependency_labels.clamp_min(0)
+            dep_indices = labels.dependency_labels.clamp_min(0) # batch x seq x label_count
         # Otherwise - check if he wants the MST or just the argmax
         elif compute_mst:
-            dep_indices = compute_mst_tree(attention_scores, extended_attention_mask)
+            dep_indices = compute_mst_tree(attention_scores[:, self.func_label_idx, :, :], extended_attention_mask[:, self.func_label_idx, :, :]).unsqueeze(-1) # batch x seq x 1
         else: 
-            dep_indices = torch.argmax(attention_scores, dim=-1)
-        
-        # After we retrieved the dependency indicies, create a tensor of teh batch indices, and and retrieve the vectors of the heads to calculate the function
-        batch_indices = torch.arange(dep_indices.size(0)).view(-1, 1).expand(-1, dep_indices.size(1)).to(dep_indices.device)
-        dep_vectors = hidden_states[batch_indices, dep_indices, :] # batch x seq x dim
-
-        # concatenate that with the last hidden states, and send to the classifier output
-        cls_inputs = torch.cat((hidden_states, dep_vectors), dim=-1)
-        function_logits = self.cls(cls_inputs)
+            dep_indices = torch.argmax(attention_scores, dim=-1).transpose(-1, -2) # batch x seq x label_count
+            
+        function_logits = None
+        if self.cls:
+            # After we retrieved the dependency indicies, create a tensor of teh batch indices, and and retrieve the vectors of the heads to calculate the function
+            # Equivalent to:         
+            # batch_indices = torch.arange(dep_indices.size(0)).view(-1, 1).expand(-1, dep_indices.size(1)).to(dep_indices.device)
+            # hidden_states[batch_indices, dep_indices, :] # batch x seq x dim
+            dep_vectors = torch.gather(hidden_states, 1, 
+                            dep_indices[:, :, self.func_label_idx].unsqueeze(-1).expand(-1, -1, hidden_states.size(-1)))
+
+            # concatenate that with the last hidden states, and send to the classifier output
+            cls_inputs = torch.cat((hidden_states, dep_vectors), dim=-1)
+            function_logits = self.cls(cls_inputs)
 
         loss = None
         if labels is not None:
@@ -91,19 +221,72 @@ class BertSyntaxParsingHead(nn.Module):
             # step 1: dependency scores loss - this is applied to the attention scores
             loss = loss_fct(attention_scores.view(-1, hidden_states.size(-2)), labels.dependency_labels.view(-1))
             # step 2: function loss
-            loss += loss_fct(function_logits.view(-1, self.num_function_classes), labels.function_labels.view(-1))
+            if self.cls:
+                loss += loss_fct(function_logits.view(-1, self.num_function_classes), labels.function_labels.view(-1))
         
         return (loss, SyntaxLogitsOutput(attention_scores, function_logits, dep_indices))
 
+def can_func_take_parameter(fn, param_name):
+    signature = inspect.signature(fn)
+    # Exclude 'self' from parameters
+    parameters = [p.name for p in signature.parameters.values() if p.name != 'self']
+    return 'kwargs' in parameters or param_name in parameters
 
-class BertForSyntaxParsing(BertPreTrainedModel):
-
-    def __init__(self, config):
+class BertLayerWrapper(nn.Module):
+    def __init__(self, layer_idx: int, bert_layer: nn.Module, partial_info: nn.Module):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.bert_layer = bert_layer
+        self.partial_info = partial_info
+
+    def forward(self, **kwargs):
+        hidden_states = self.bert_layer(**kwargs)
+        hidden_states = self.partial_info(hidden_states)
+        return hidden_states
+
+def wrap_bert_layer_with_partial_info(layer_idx: int, bert_layer: nn.Module, partial_info: nn.Module):
+    if layer_idx != 8: return
+    orig_forward = bert_layer.forward
+    def new_forward(*args, **kwargs):
+        hidden_states = orig_forward(*args, **kwargs)
+        # Read partial_labels from the module instead of kwargs
+        hidden_states = partial_info(hidden_states)
+        return hidden_states
+    bert_layer.forward = new_forward
+
+class BaseForSyntaxParsing(BertPreTrainedModel):
+    base_model_prefix = ""
+    def __init__(self, config, syntax_head_size=128, bert_cls=BertModel, is_partial_info_model=False, is_cls_model=False, syntax_attn_label_count=1, syntax_func_label_idx=0):
         super().__init__(config)
-
-        self.bert = BertModel(config, add_pooling_layer=False)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
-        self.syntax = BertSyntaxParsingHead(config)
+        # conversions
+        setattr(config, "hidden_dropout_prob", getattr(config, "hidden_dropout_prob", 0.1))
+        setattr(config, "initializer_range", getattr(config, "classifier_init_range", getattr(config, 'decoder_init_range', 0.02)))
+        
+        if not hasattr(config, 'syntax_head_size'):
+            config.syntax_head_size = syntax_head_size
+        if not hasattr(config, 'syntax_attn_label_count'):
+            config.syntax_attn_label_count = syntax_attn_label_count
+        if not hasattr(config, 'syntax_func_label_idx'):
+            config.syntax_func_label_idx = syntax_func_label_idx
+        if not hasattr(config, 'is_partial_info_model'):
+            config.is_partial_info_model = is_partial_info_model
+        if not hasattr(config, 'is_cls_model'):
+            config.is_cls_model = is_cls_model
+            # if is_cls_model:
+            #    assert config.is_partial_info_model, "CLS model requires partial info model to be set to True"
+
+        self.bert = bert_cls(config, **({} if not can_func_take_parameter(bert_cls.__init__, 'add_pooling_layer') else {'add_pooling_layer': False}))
+        self.send_token_type_ids = can_func_take_parameter(self.bert.forward, 'token_type_ids')
+        self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1))
+        if config.is_partial_info_model:
+            self.partial_info = BertSyntaxPartialInfoHead(config)
+            for layer_idx, layer in enumerate(self.bert.encoder.layer):
+                wrap_bert_layer_with_partial_info(layer_idx, layer, self.partial_info)
+            # self.bert.encoder.layer = nn.ModuleList([BertLayerWrapper(layer_idx, layer, self.partial_info) for layer_idx, layer in enumerate(self.bert.encoder.layer)])
+        if config.is_cls_model:
+            self.syntax = BertSyntaxValidClassifierHead(config)
+        else:
+            self.syntax = BertSyntaxParsingHead(config)
         
         # Initialize weights and apply final processing
         self.post_init()
@@ -114,7 +297,8 @@ class BertForSyntaxParsing(BertPreTrainedModel):
         attention_mask: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
-        labels: Optional[SyntaxLabels] = None,
+        partial_labels: Optional[torch.Tensor] = None,
+        labels: Optional[Union[SyntaxLabels, torch.Tensor]] = None,
         head_mask: Optional[torch.Tensor] = None,
         inputs_embeds: Optional[torch.Tensor] = None,
         output_attentions: Optional[bool] = None,
@@ -124,25 +308,49 @@ class BertForSyntaxParsing(BertPreTrainedModel):
     ):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        if not self.config.is_partial_info_model and partial_labels:
+            raise ValueError('Cannot pass partial label when model not initialized with partial info')
+
+        kwargs = dict(token_type_ids=token_type_ids, head_mask=head_mask) if self.send_token_type_ids else {}
+        if self.config.is_partial_info_model:
+            # Store partial_labels on the module (avoids passing through kwargs which newer BERT doesn't support)
+            self.partial_info.set_partial_labels(partial_labels)
+            # Clear auxiliary losses before forward pass
+            self.partial_info.clear_aux_losses()
+
         bert_outputs = self.bert(
             input_ids,
             attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
             position_ids=position_ids,
-            head_mask=head_mask,
             inputs_embeds=inputs_embeds,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            **kwargs
         )
         
         extended_attention_mask = None
         if attention_mask is not None:
             extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.size())
+        
+        hidden_states = self.dropout(bert_outputs[0])
+
+        # if self.config.is_partial_info_model:
+        #     partial_labels = partial_labels or SyntaxLabels(dependency_labels=torch.full_like(input_ids, -1), function_labels=torch.full_like(input_ids, -1))
+        #     hidden_states = self.partial_info(hidden_states, partial_labels)
+
         # apply the syntax head
-        loss, logits = self.syntax(self.dropout(bert_outputs[0]), extended_attention_mask, labels, compute_syntax_mst)
+        loss, logits = self.syntax(hidden_states, extended_attention_mask, labels, compute_syntax_mst)
         
+        # Add auxiliary loss from partial info head (weighted lightly)
+        if self.config.is_partial_info_model and self.training:
+            aux_loss = self.partial_info.get_total_aux_loss()
+            if aux_loss is not None and loss is not None:
+                loss = loss + aux_loss
+
         if not return_dict:
+            if self.config.is_cls_model:
+                return (loss, logits) + bert_outputs[2:]
             return (loss,(logits.dependency_logits, logits.function_logits)) + bert_outputs[2:]
         
         return SyntaxTaggingOutput(
@@ -152,7 +360,15 @@ class BertForSyntaxParsing(BertPreTrainedModel):
             attentions=bert_outputs.attentions,
         )
 
+    def get_input_embeddings(self):
+        return self.bert.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.bert.embeddings.word_embeddings = value
+
     def predict(self, sentences: Union[str, List[str]], tokenizer: BertTokenizerFast, compute_mst=True):
+        if self.config.is_cls_model:
+            raise ValueError('Cannot use predict function with classification model')
         if isinstance(sentences, str):
             sentences = [sentences]
             
@@ -160,17 +376,32 @@ class BertForSyntaxParsing(BertPreTrainedModel):
         inputs = tokenizer(sentences, padding='longest', truncation=True, return_tensors='pt')
         inputs = {k:v.to(self.device) for k,v in inputs.items()}
         logits = self.forward(**inputs, return_dict=True, compute_syntax_mst=compute_mst).logits
-        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits)
+        return parse_logits(inputs['input_ids'].tolist(), sentences, tokenizer, logits, self.syntax.func_label_idx)
+
+class AutoForSyntaxParsing(_BaseAutoModelClass):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *args, **kwargs):
+        auto_cfg = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
+        custom_kwargs = dict(syntax_head_size=128, is_partial_info_model=False, is_cls_model=False, syntax_attn_label_count=1, syntax_func_label_idx=0)
+        custom_kwargs = {k:kwargs.pop(k, getattr(auto_cfg, k, v)) for k,v in custom_kwargs.items()}
+        base_cls = BaseForSyntaxParsing
+        with no_init_weights():
+            bert_cls = AutoModel.from_config(auto_cfg, *args, **kwargs).__class__
+            if 'Syntax' in bert_cls.__name__:
+                base_cls = bert_cls
+                
+        return base_cls.from_pretrained(pretrained_model_name_or_path, *args, **kwargs, **custom_kwargs, bert_cls=bert_cls, key_mapping={"^model": "bert"})
+    
 
-def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: SyntaxLogitsOutput):
+def parse_logits(input_ids: List[List[int]], sentences: List[str], tokenizer: BertTokenizerFast, logits: SyntaxLogitsOutput, func_label_idx: int = 0):
     outputs = []        
 
     special_toks = tokenizer.all_special_tokens
     special_toks.remove(tokenizer.unk_token)
     special_toks.remove(tokenizer.mask_token)
-    
     for i in range(len(sentences)):
-        deps = logits.dependency_head_indices[i].tolist()
+        # dependency_head_indices is seq x label_count - the tree is built from the function label's column
+        deps = logits.dependency_head_indices[i][:, func_label_idx].tolist()
         funcs = logits.function_logits.argmax(-1)[i].tolist()
         toks = [tok for tok in tokenizer.convert_ids_to_tokens(input_ids[i]) if tok not in special_toks]
         
@@ -312,4 +543,7 @@ def choose_contracting_arc(indices: torch.LongTensor, sorted_indices: List[List[
     if best_base_idx == -1:
         raise ValueError('Stuck in endless loop trying to compute syntax mst. Please try again setting compute_syntax_mst=False')
 
-    return best_base_idx, best_head_idx
+    return best_base_idx, best_head_idx
+
+def replace_tensor_value(tensor, orig_value, replacement_value):
+    return torch.where(tensor == orig_value, replacement_value, tensor)