Upload ConlluTokenClassificationPipeline

encoder.py

@@ -2,7 +2,7 @@ import torch
 from torch import nn
 from torch import Tensor, LongTensor
 
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 try:
     from peft import LoraConfig, get_peft_model
@@ -28,17 +28,13 @@ class WordTransformerEncoder(nn.Module):
     ):
         super().__init__()
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModel.from_pretrained(model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
         if use_lora:
             if not PEFT_AVAILABLE:
                 raise ImportError("peft is required for LoRA fine-tuning. Install with `pip install peft`.")
             if lora_target_modules is None:
-                # XLM-RoBERTa и Roberta-family
-                if "roberta" in model_name.lower():
-                    lora_target_modules = ["q_proj", "v_proj"]
-                else:
-                    lora_target_modules = ["query", "value"]
+                  lora_target_modules = ["query", "value"]
             lora_config = LoraConfig(
                 r=lora_r,
                 lora_alpha=lora_alpha,
@@ -88,7 +84,8 @@ class WordTransformerEncoder(nn.Module):
         ])
 
         # Run model and extract subtokens embeddings from the last layer.
-        subtokens_embeddings = self.model(**subtokens).last_hidden_state
+        outputs = self.model(**subtokens, output_hidden_states=True)
+        subtokens_embeddings = outputs.hidden_states[-1] 
 
         # Aggreate subtokens embeddings into words embeddings.
         # [batch_size, n_words, embedding_size]
@@ -137,7 +134,7 @@ class WordTransformerEncoder(nn.Module):
 
     def get_embeddings_layer(self):
         """Returns the embeddings model."""
-        return self.model.embeddings
+        return self.model.roberta.embeddings
         
     def get_transformer_layers(self) -> list[nn.Module]:
         """

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:37137a12604aa1ee98f7cc4627a3ed76c63cdb9eb9d2fd02c4a73aaf17325dae
-size 1134190536
+oid sha256:b27ba6b77caece9754d522ddc5bf9e63844a10a04e972b65961a9baccbb08bf5
+size 1134198480