fix(model): correctly configure AutoModelForTokenClassification

config.json

@@ -5,7 +5,8 @@
   "auto_map": {
     "AutoConfig": "configuration_nllbllm2vec.NLLBLLM2VecConfig",
     "AutoModel": "modeling_nllbllm2vec.NLLBLLM2Vec",
-    "AutoModelForSequenceClassification": "modeling_nllbllm2vec.NLLBLLM2VecForSequenceClassification"
+    "AutoModelForSequenceClassification": "modeling_nllbllm2vec.NLLBLLM2VecForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_nllbllm2vec.NLLBLLM2VecForTokenClassification"
   },
   "llm2vec_config": {
     "_name_or_path": "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp",

modeling_nllbllm2vec.py

@@ -18,7 +18,7 @@ from transformers.modeling_outputs import (
     TokenClassifierOutput,
 )
 from transformers.modeling_utils import PreTrainedModel
-from transformers.models.auto import AutoModel, AutoModelForSequenceClassification
+from transformers.models.auto import AutoModel, AutoModelForSequenceClassification, AutoModelForTokenClassification
 from transformers.models.m2m_100.modeling_m2m_100 import M2M100Encoder
 from transformers.tokenization_utils import BatchEncoding
 
@@ -536,6 +536,6 @@ AutoModel.register(NLLBLLM2VecConfig, NLLBLLM2Vec)
 AutoModelForSequenceClassification.register(
     NLLBLLM2VecConfig, NLLBLLM2VecForSequenceClassification
 )
-AutoModelForSequenceClassification.register(
+AutoModelForTokenClassification.register(
     NLLBLLM2VecConfig, NLLBLLM2VecForTokenClassification
 )