Update modeling.py

modeling.py

@@ -3,7 +3,6 @@ from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
-from torchcrf import CRF
 from transformers import PretrainedConfig, PreTrainedModel
 from transformers.modeling_outputs import TokenClassifierOutput
 
@@ -16,8 +15,6 @@ except Exception:
         EuroBertModel = None
         print("COULD NOT IMPORT EUROBERT MODEL")
 
-# Large negative number for masking impossible transitions
-LARGE_NEGATIVE_NUMBER = -1e9
 NUM_PER_LAYER = 16
 
 
@@ -64,727 +61,11 @@ def _build_backbone_from_config(config):
     return backbone, backbone_name
 
 
-class MultiHeadCRFConfig(PretrainedConfig):
-    """
-    Configuration class for Multi-Head CRF models.
-    """
-
-    model_type = "multihead-crf-tagger"
-
-    def __init__(
-        self,
-        entity_types: Optional[List[str]] = None,
-        number_of_layers_per_head: int = 1,
-        crf_reduction: str = "mean",
-        freeze_backbone: bool = False,
-        num_frozen_encoders: int = 0,
-        classifier_dropout: float = 0.1,
-        classifier_hidden_layers: Optional[Tuple] = None,
-        class_weights: Optional[List[float]] = None,
-        backbone_model_name: Optional[str] = None,
-        **kwargs,
-    ):
-        self.entity_types = entity_types or []
-        self.number_of_layers_per_head = number_of_layers_per_head
-        self.crf_reduction = crf_reduction
-        self.freeze_backbone = freeze_backbone
-        self.num_frozen_encoders = num_frozen_encoders
-        self.classifier_dropout = classifier_dropout
-        self.classifier_hidden_layers = classifier_hidden_layers
-        self.class_weights = class_weights
-        self.backbone_model_name = backbone_model_name
-        super().__init__(**kwargs)
-
-
-class MultiHeadCRF(nn.Module):
-    """
-    Custom CRF implementation with BIO transition masking.
-    """
-
-    def __init__(self, num_tags: int, batch_first: bool = True) -> None:
-        if num_tags <= 0:
-            raise ValueError(f"invalid number of tags: {num_tags}")
-        super().__init__()
-        self.num_tags = num_tags
-        self.batch_first = batch_first
-        self.start_transitions = nn.Parameter(torch.empty(num_tags))
-        self.end_transitions = nn.Parameter(torch.empty(num_tags))
-        self.transitions = nn.Parameter(torch.empty(num_tags, num_tags))
-
-        self.reset_parameters()
-        self.mask_impossible_transitions()
-
-    def reset_parameters(self) -> None:
-        nn.init.uniform_(self.start_transitions, -0.1, 0.1)
-        nn.init.uniform_(self.end_transitions, -0.1, 0.1)
-        nn.init.uniform_(self.transitions, -0.1, 0.1)
-
-    def mask_impossible_transitions(self) -> None:
-        with torch.no_grad():
-            if self.num_tags > 2:
-                self.start_transitions[2] = LARGE_NEGATIVE_NUMBER
-                self.transitions[0][2] = LARGE_NEGATIVE_NUMBER
-
-            if self.num_tags > 3:
-                self.start_transitions[3] = LARGE_NEGATIVE_NUMBER
-                for i in range(3):
-                    self.transitions[i][3] = LARGE_NEGATIVE_NUMBER
-                for i in range(3):
-                    self.transitions[3][i] = LARGE_NEGATIVE_NUMBER
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(num_tags={self.num_tags})"
-
-    def forward(
-        self,
-        emissions: torch.Tensor,
-        tags: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-        reduction: str = "mean",
-    ) -> torch.Tensor:
-        self._validate(emissions, tags=tags, mask=mask)
-        if reduction not in ("none", "sum", "mean", "token_mean"):
-            raise ValueError(f"invalid reduction: {reduction}")
-        if mask is None:
-            mask = torch.ones_like(tags, dtype=torch.uint8)
-
-        device = emissions.device
-        tags = tags.to(device)
-        mask = mask.to(device)
-
-        if self.batch_first:
-            emissions = emissions.transpose(0, 1)
-            tags = tags.transpose(0, 1)
-            mask = mask.transpose(0, 1)
-
-        numerator = self._compute_score(emissions, tags, mask)
-        denominator = self._compute_normalizer(emissions, mask)
-        llh = numerator - denominator
-        nllh = -llh
-
-        if reduction == "none":
-            return nllh
-        if reduction == "sum":
-            return nllh.sum()
-        if reduction == "mean":
-            return nllh.mean()
-        return nllh.sum() / mask.type_as(emissions).sum()
-
-    def decode(
-        self, emissions: torch.Tensor, mask: Optional[torch.Tensor] = None
-    ) -> List[List[int]]:
-        self._validate(emissions, mask=mask)
-        if mask is None:
-            mask = emissions.new_ones(emissions.shape[:2], dtype=torch.uint8)
-
-        if self.batch_first:
-            emissions = emissions.transpose(0, 1)
-            mask = mask.transpose(0, 1)
-
-        return self._viterbi_decode(emissions, mask)
-
-    def _validate(
-        self,
-        emissions: torch.Tensor,
-        tags: Optional[torch.Tensor] = None,
-        mask: Optional[torch.Tensor] = None,
-    ) -> None:
-        if emissions.dim() != 3:
-            raise ValueError(
-                f"emissions must have dimension of 3, got {emissions.dim()}"
-            )
-        if emissions.size(2) != self.num_tags:
-            raise ValueError(
-                f"expected last dimension of emissions is {self.num_tags}, "
-                f"got {emissions.size(2)}"
-            )
-
-        if tags is not None and emissions.shape[:2] != tags.shape:
-            raise ValueError(
-                "the first two dimensions of emissions and tags must match, "
-                f"got {tuple(emissions.shape[:2])} and {tuple(tags.shape)}"
-            )
-
-        if mask is not None:
-            if emissions.shape[:2] != mask.shape:
-                raise ValueError(
-                    "the first two dimensions of emissions and mask must match, "
-                    f"got {tuple(emissions.shape[:2])} and {tuple(mask.shape)}"
-                )
-            no_empty_seq = not self.batch_first and mask[0].all()
-            no_empty_seq_bf = self.batch_first and mask[:, 0].all()
-            if not no_empty_seq and not no_empty_seq_bf:
-                raise ValueError("mask of the first timestep must all be on")
-
-    def _compute_score(
-        self, emissions: torch.Tensor, tags: torch.Tensor, mask: torch.Tensor
-    ) -> torch.Tensor:
-        assert emissions.dim() == 3 and tags.dim() == 2
-        assert emissions.shape[:2] == tags.shape
-        assert emissions.size(2) == self.num_tags
-        assert mask.shape == tags.shape
-        assert mask[0].all()
-
-        device = emissions.device
-        tags = tags.to(device)
-        mask = mask.to(device)
-
-        seq_length, batch_size = tags.shape
-        mask = mask.type_as(emissions)
-
-        batch_indices = torch.arange(batch_size, device=device)
-        score = self.start_transitions[tags[0]]
-        score += emissions[0, batch_indices, tags[0]]
-
-        for i in range(1, seq_length):
-            score += self.transitions[tags[i - 1], tags[i]] * mask[i]
-            score += emissions[i, batch_indices, tags[i]] * mask[i]
-
-        seq_ends = mask.long().sum(dim=0) - 1
-        last_tags = tags[seq_ends, batch_indices]
-        score += self.end_transitions[last_tags]
-
-        return score
-
-    def _compute_normalizer(
-        self, emissions: torch.Tensor, mask: torch.Tensor
-    ) -> torch.Tensor:
-        assert emissions.dim() == 3 and mask.dim() == 2
-        assert emissions.shape[:2] == mask.shape
-        assert emissions.size(2) == self.num_tags
-        assert mask[0].all()
-
-        seq_length = emissions.size(0)
-        score = self.start_transitions + emissions[0]
-
-        for i in range(1, seq_length):
-            broadcast_score = score.unsqueeze(2)
-            broadcast_emissions = emissions[i].unsqueeze(1)
-            next_score = broadcast_score + self.transitions + broadcast_emissions
-            next_score = torch.logsumexp(next_score, dim=1)
-            score = torch.where(mask[i].unsqueeze(1).bool(), next_score, score)
-
-        score += self.end_transitions
-        return torch.logsumexp(score, dim=1)
-
-    def _viterbi_decode(
-        self, emissions: torch.Tensor, mask: torch.Tensor
-    ) -> List[List[int]]:
-        assert emissions.dim() == 3 and mask.dim() == 2
-        assert emissions.shape[:2] == mask.shape
-        assert emissions.size(2) == self.num_tags
-        assert mask[0].all()
-
-        seq_length, batch_size = mask.shape
-        score = self.start_transitions + emissions[0]
-        history = []
-
-        for i in range(1, seq_length):
-            broadcast_score = score.unsqueeze(2)
-            broadcast_emission = emissions[i].unsqueeze(1)
-            next_score = broadcast_score + self.transitions + broadcast_emission
-            next_score, indices = next_score.max(dim=1)
-            score = torch.where(mask[i].unsqueeze(1).bool(), next_score, score)
-            history.append(indices)
-
-        score += self.end_transitions
-
-        seq_ends = mask.long().sum(dim=0) - 1
-        best_tags_list = []
-
-        for idx in range(batch_size):
-            _, best_last_tag = score[idx].max(dim=0)
-            best_tags = [best_last_tag.item()]
-
-            for hist in reversed(history[: seq_ends[idx]]):
-                best_last_tag = hist[idx][best_tags[-1]]
-                best_tags.append(best_last_tag.item())
-
-            best_tags.reverse()
-            best_tags_list.append(best_tags)
-
-        return best_tags_list
-
-
-class TokenClassificationModelCRF(PreTrainedModel):
-    """
-    Custom token classification model with CRF layer and configurable classifier head.
-    """
-
-    def __init__(
-        self,
-        config,
-        base_model=None,
-        freeze_backbone=False,
-        classifier_hidden_layers=None,
-        classifier_dropout=0.1,
-    ):
-        super().__init__(config)
-        self.config = config
-        self.num_labels = config.num_labels
-
-        if base_model is None:
-            self.roberta, backbone_name = _build_backbone_from_config(config)
-        else:
-            if hasattr(base_model, "roberta"):
-                self.roberta = base_model.roberta
-            else:
-                self.roberta = base_model
-            backbone_name = (
-                getattr(getattr(self.roberta, "config", None), "_name_or_path", None)
-                or getattr(config, "backbone_model_name", None)
-                or getattr(config, "_name_or_path", None)
-            )
-            if getattr(config, "backbone_model_name", None) is None:
-                config.backbone_model_name = backbone_name
-
-        self.lm_output_size = self.roberta.config.hidden_size
-
-        self.config.freeze_backbone = freeze_backbone
-        self.config.classifier_hidden_layers = classifier_hidden_layers
-        self.config.classifier_dropout = classifier_dropout
-
-        if freeze_backbone:
-            print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n")
-            for param in self.roberta.parameters():
-                param.requires_grad = False
-            self.roberta.eval()
-        else:
-            print("+" * 30, "\n\n", "NOT Freezing backbone...", "+" * 30, "\n\n")
-            self.roberta.train(True)
-
-        self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1))
-        self.crf = CRF(self.num_labels, batch_first=True)
-
-        self._build_classifier_head(classifier_hidden_layers, classifier_dropout)
-        self.post_init()
-
-    def _build_classifier_head(self, hidden_layers, dropout_rate):
-        layers = []
-        input_size = self.lm_output_size
-
-        if not hidden_layers:
-            self.classifier = nn.Sequential(
-                nn.Dropout(dropout_rate), nn.Linear(input_size, self.num_labels)
-            )
-            return
-
-        for hidden_size in hidden_layers:
-            layers.append(nn.Linear(input_size, hidden_size))
-            layers.append(nn.ReLU())
-            layers.append(nn.Dropout(dropout_rate))
-            input_size = hidden_size
-
-        layers.append(nn.Linear(input_size, self.num_labels))
-        self.classifier = nn.Sequential(*layers)
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs,
-    ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-
-        try:
-            outputs = self.roberta(
-                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,
-            )
-        except TypeError:
-            outputs = self.roberta(
-                input_ids,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                inputs_embeds=inputs_embeds,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
-
-        sequence_output = self.dropout(outputs.last_hidden_state)
-        logits = self.classifier(sequence_output)
-
-        loss = None
-        if labels is not None:
-            labels_long = labels.long()
-            if attention_mask is not None:
-                mask = attention_mask.bool()
-                loss = -self.crf(logits, labels_long, mask=mask, reduction="mean")
-            else:
-                if not getattr(self, "_warned_no_attention_mask", False):
-                    print(
-                        "WARNING: attention_mask is None; CRF loss will include padding tokens."
-                    )
-                    self._warned_no_attention_mask = True
-                loss = -self.crf(logits, labels_long, reduction="mean")
-
-        if not return_dict:
-            output = (logits,) + outputs[2:]
-            return ((loss,) + output) if loss is not None else output
-
-        return TokenClassifierOutput(
-            loss=loss,
-            logits=logits,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    @property
-    def device_info(self):
-        return next(self.parameters()).device
-
-    def get_input_embeddings(self):
-        return self.roberta.get_input_embeddings()
-
-    def set_input_embeddings(self, value):
-        self.roberta.set_input_embeddings(value)
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        config = kwargs.pop("config", None)
-        if config is None:
-            from transformers import AutoConfig
-
-            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
-
-        freeze_backbone = getattr(config, "freeze_backbone", False)
-        classifier_hidden_layers = getattr(config, "classifier_hidden_layers", None)
-        classifier_dropout = getattr(config, "classifier_dropout", 0.1)
-
-        model = cls(
-            config=config,
-            freeze_backbone=freeze_backbone,
-            classifier_hidden_layers=classifier_hidden_layers,
-            classifier_dropout=classifier_dropout,
-        )
-
-        try:
-            state_dict = torch.load(
-                f"{pretrained_model_name_or_path}/pytorch_model.bin", map_location="cpu"
-            )
-            model.load_state_dict(state_dict)
-        except Exception:
-            print(
-                "Warning: Could not load pre-trained weights. Using randomly initialized model."
-            )
-
-        return model
-
-
-class TokenClassificationModelMultiHeadCRF(PreTrainedModel):
-    """
-    Multi-Head CRF model for token classification with multiple entity types.
-    """
-
-    config_class = MultiHeadCRFConfig
-    base_model_prefix = "roberta"
-    _keys_to_ignore_on_load_unexpected = [r"pooler"]
-
-    def __init__(self, config, base_model=None, freeze_backbone=None):
-        super().__init__(config)
-        self.config = config
-
-        self.entity_types = getattr(config, "entity_types", [])
-        if not self.entity_types:
-            raise ValueError("entity_types must be provided in config")
-
-        self.num_labels = config.num_labels
-        self.number_of_layers_per_head = getattr(config, "number_of_layers_per_head", 1)
-        self.crf_reduction = getattr(config, "crf_reduction", "mean")
-        freeze_backbone = (
-            freeze_backbone
-            if freeze_backbone is not None
-            else getattr(config, "freeze_backbone", False)
-        )
-        self.num_frozen_encoders = getattr(config, "num_frozen_encoders", 0)
-        classifier_dropout = getattr(config, "classifier_dropout", 0.1)
-
-        if base_model is None:
-            self.roberta, backbone_name = _build_backbone_from_config(config)
-        else:
-            if hasattr(base_model, "roberta"):
-                self.roberta = base_model.roberta
-            else:
-                self.roberta = base_model
-            backbone_name = (
-                getattr(getattr(self.roberta, "config", None), "_name_or_path", None)
-                or getattr(config, "backbone_model_name", None)
-                or getattr(config, "_name_or_path", None)
-            )
-            if getattr(config, "backbone_model_name", None) is None:
-                config.backbone_model_name = backbone_name
-
-        self.hidden_size = self.roberta.config.hidden_size
-        self.dropout = nn.Dropout(getattr(config, "hidden_dropout_prob", 0.1))
-
-        print(f"Creating Multi-Head CRF with entity types: {sorted(self.entity_types)}")
-
-        for entity_type in self.entity_types:
-            for i in range(self.number_of_layers_per_head):
-                setattr(
-                    self,
-                    f"{entity_type}_dense_{i}",
-                    nn.Linear(self.hidden_size, self.hidden_size),
-                )
-                setattr(
-                    self,
-                    f"{entity_type}_dense_activation_{i}",
-                    nn.GELU(approximate="none"),
-                )
-                setattr(
-                    self, f"{entity_type}_dropout_{i}", nn.Dropout(classifier_dropout)
-                )
-
-            setattr(
-                self,
-                f"{entity_type}_classifier",
-                nn.Linear(self.hidden_size, self.num_labels),
-            )
-            setattr(
-                self,
-                f"{entity_type}_crf",
-                MultiHeadCRF(num_tags=self.num_labels, batch_first=True),
-            )
-
-        if freeze_backbone:
-            self._freeze_backbone()
-
-        self.post_init()
-
-    def _freeze_backbone(self):
-        print("+" * 30, "\n\n", "Freezing backbone...", "+" * 30, "\n\n")
-
-        for param in self.roberta.embeddings.parameters():
-            param.requires_grad = False
-
-        if self.num_frozen_encoders > 0:
-            for _, param in islice(
-                self.roberta.encoder.named_parameters(),
-                self.num_frozen_encoders * NUM_PER_LAYER,
-            ):
-                param.requires_grad = False
-
-    def reset_head_parameters(self):
-        for entity_type in self.entity_types:
-            for i in range(self.number_of_layers_per_head):
-                getattr(self, f"{entity_type}_dense_{i}").reset_parameters()
-            getattr(self, f"{entity_type}_classifier").reset_parameters()
-            getattr(self, f"{entity_type}_crf").reset_parameters()
-            getattr(self, f"{entity_type}_crf").mask_impossible_transitions()
-
-    def forward(
-        self,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        token_type_ids: Optional[torch.LongTensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        head_mask: Optional[torch.FloatTensor] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[Dict[str, torch.LongTensor]] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **kwargs,
-    ):
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-
-        try:
-            outputs = self.roberta(
-                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,
-            )
-        except TypeError:
-            outputs = self.roberta(
-                input_ids,
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                inputs_embeds=inputs_embeds,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
-
-        sequence_output = outputs[0]
-        sequence_output = self.dropout(sequence_output)
-
-        logits = {}
-        for entity_type in self.entity_types:
-            head_output = sequence_output
-            for i in range(self.number_of_layers_per_head):
-                head_output = getattr(self, f"{entity_type}_dense_{i}")(head_output)
-                head_output = getattr(self, f"{entity_type}_dense_activation_{i}")(
-                    head_output
-                )
-                head_output = getattr(self, f"{entity_type}_dropout_{i}")(head_output)
-            logits[entity_type] = getattr(self, f"{entity_type}_classifier")(
-                head_output
-            )
-
-        if labels is not None:
-            losses = {}
-            mask = attention_mask.bool() if attention_mask is not None else None
-
-            for entity_type in self.entity_types:
-                if entity_type in labels:
-                    entity_labels = (
-                        labels[entity_type].long().to(logits[entity_type].device)
-                    )
-                    crf = getattr(self, f"{entity_type}_crf")
-                    if mask is not None:
-                        losses[entity_type] = crf(
-                            logits[entity_type],
-                            entity_labels,
-                            mask=mask,
-                            reduction=self.crf_reduction,
-                        )
-                    else:
-                        if not getattr(self, "_warned_no_attention_mask", False):
-                            print(
-                                "WARNING: attention_mask is None; CRF loss will include padding tokens."
-                            )
-                            self._warned_no_attention_mask = True
-                        losses[entity_type] = crf(
-                            logits[entity_type],
-                            entity_labels,
-                            reduction=self.crf_reduction,
-                        )
-
-            total_loss = sum(losses.values())
-            return total_loss, logits
-
-        predictions = {}
-        mask = attention_mask.bool() if attention_mask is not None else None
-
-        for entity_type in self.entity_types:
-            crf = getattr(self, f"{entity_type}_crf")
-            if mask is not None:
-                decoded = crf.decode(logits[entity_type], mask=mask)
-            else:
-                decoded = crf.decode(logits[entity_type])
-            predictions[entity_type] = torch.tensor(decoded)
-
-        return [predictions[ent] for ent in sorted(self.entity_types)]
-
-    def get_input_embeddings(self):
-        return self.roberta.get_input_embeddings()
-
-    def set_input_embeddings(self, value):
-        self.roberta.set_input_embeddings(value)
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        import json
-        import os
-
-        config = kwargs.pop("config", None)
-
-        if config is None:
-            config_file = os.path.join(pretrained_model_name_or_path, "config.json")
-            if os.path.exists(config_file):
-                with open(config_file, "r") as f:
-                    config_dict = json.load(f)
-                config = MultiHeadCRFConfig(**config_dict)
-            else:
-                from transformers import AutoConfig
-
-                config = AutoConfig.from_pretrained(
-                    pretrained_model_name_or_path,
-                    trust_remote_code=kwargs.get("trust_remote_code", True),
-                )
-
-        roberta_defaults = {
-            "layer_norm_eps": 1e-5,
-            "hidden_size": 768,
-            "num_hidden_layers": 12,
-            "num_attention_heads": 12,
-            "intermediate_size": 3072,
-            "hidden_act": "gelu",
-            "hidden_dropout_prob": 0.1,
-            "attention_probs_dropout_prob": 0.1,
-            "max_position_embeddings": 514,
-            "type_vocab_size": 1,
-            "initializer_range": 0.02,
-            "vocab_size": 52000,
-            "pad_token_id": 1,
-            "bos_token_id": 0,
-            "eos_token_id": 2,
-            "position_embedding_type": "absolute",
-            "use_cache": True,
-            "is_decoder": False,
-            "add_cross_attention": False,
-            "chunk_size_feed_forward": 0,
-            "output_hidden_states": False,
-            "output_attentions": False,
-            "torchscript": False,
-            "tie_word_embeddings": True,
-            "return_dict": True,
-            "gradient_checkpointing": False,
-            "pruned_heads": {},
-            "problem_type": None,
-            "embedding_size": None,
-        }
-
-        for key, default_value in roberta_defaults.items():
-            if not hasattr(config, key) or getattr(config, key) is None:
-                setattr(config, key, default_value)
-
-        freeze_backbone = getattr(config, "freeze_backbone", False)
-        model = cls(config=config, freeze_backbone=freeze_backbone)
-
-        weight_file = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
-        safetensors_file = os.path.join(
-            pretrained_model_name_or_path, "model.safetensors"
-        )
-
-        try:
-            if os.path.exists(safetensors_file):
-                from safetensors.torch import load_file
-
-                state_dict = load_file(safetensors_file)
-                model.load_state_dict(state_dict)
-            elif os.path.exists(weight_file):
-                state_dict = torch.load(weight_file, map_location="cpu")
-                model.load_state_dict(state_dict)
-            else:
-                print(
-                    "Warning: No pre-trained weights found. Using randomly initialized model."
-                )
-        except Exception as e:
-            print(f"Warning: Could not load pre-trained weights: {e}")
-
-        return model
 
 
 class MultiHeadConfig(PretrainedConfig):
     """
-    Configuration class for Multi-Head models (without CRF).
+    Configuration class for Multi-Head models.
     """
 
     model_type = "multihead-tagger"
@@ -814,7 +95,7 @@ class MultiHeadConfig(PretrainedConfig):
 
 class TokenClassificationModelMultiHead(PreTrainedModel):
     """
-    Multi-Head model for token classification with multiple entity types (no CRF).
+    Multi-Head model for token classification with multiple entity types.
     """
 
     config_class = MultiHeadConfig
@@ -1104,7 +385,7 @@ class TokenClassificationModelMultiHead(PreTrainedModel):
 
 class TokenClassificationModel(PreTrainedModel):
     """
-    Custom token classification model with configurable classifier head (no CRF).
+    Custom token classification model with configurable classifier head.
     """
 
     def __init__(self, config, base_model=None):
@@ -1272,46 +553,6 @@ def load_custom_cardioner_multiclass_model(model_path: str, device: str = "auto"
     return model, tokenizer, model.config
 
 
-def load_custom_multihead_crf_model(model_path: str, device: str = "auto"):
-    import json
-    import os
-
-    from transformers import AutoTokenizer
-
-    required_files = ["config.json", "modeling.py"]
-    missing_files = [
-        f for f in required_files if not os.path.exists(os.path.join(model_path, f))
-    ]
-
-    if missing_files:
-        raise FileNotFoundError(
-            f"Missing required files in {model_path}: {missing_files}"
-        )
-
-    print(f"Loading Multi-Head CRF model from: {model_path}")
-
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-
-    with open(os.path.join(model_path, "config.json"), "r") as f:
-        config_dict = json.load(f)
-
-    config = MultiHeadCRFConfig(**config_dict)
-
-    model = TokenClassificationModelMultiHeadCRF.from_pretrained(
-        model_path, config=config
-    )
-
-    if device == "auto":
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    model = model.to(device)
-
-    print(f"Model loaded successfully on {device}")
-    print(f"Model type: {type(model).__name__}")
-    print(f"Entity types: {model.entity_types}")
-    print(f"Number of labels per head: {model.num_labels}")
-
-    return model, tokenizer, model.config
 
 
 def validate_custom_multiclass_model_directory(model_path: str) -> dict:
@@ -1376,9 +617,6 @@ def validate_custom_multiclass_model_directory(model_path: str) -> dict:
             validation_results["model_info"]["freeze_backbone"] = config.get(
                 "freeze_backbone", None
             )
-            validation_results["model_info"]["use_crf"] = (
-                "TokenClassificationModelCRF" in str(config.get("architectures", []))
-            )
 
             if not config.get("auto_map"):
                 validation_results["warnings"].append(
@@ -1397,7 +635,6 @@ def validate_custom_multiclass_model_directory(model_path: str) -> dict:
 
             required_classes = [
                 "TokenClassificationModel",
-                "TokenClassificationModelCRF",
             ]
             missing_classes = [cls for cls in required_classes if cls not in content]
 
@@ -1415,12 +652,6 @@ def validate_custom_multiclass_model_directory(model_path: str) -> dict:
     return validation_results
 
 
-try:
-    from transformers import AutoConfig
-
-    AutoConfig.register("multihead-crf-tagger", MultiHeadCRFConfig)
-except Exception:
-    pass
 
 
 def patch_legacy_model(
@@ -1492,4 +723,4 @@ def patch_multiple_models(
         f"Successfully {'would patch' if dry_run else 'patched'}: {success}/{len(model_paths)}"
     )
 
-    return results
+    return results