Upload ConlluTokenClassificationPipeline

encoder.py

@@ -2,7 +2,7 @@ import torch
 from torch import nn
 from torch import Tensor, LongTensor
 
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModel
 
 try:
     from peft import LoraConfig, get_peft_model
@@ -28,30 +28,23 @@ class WordTransformerEncoder(nn.Module):
     ):
         super().__init__()
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        self.model = AutoModel.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:
-                  lora_target_modules = ["query", "value"]
+              lora_target_modules = ["query", "value"]
             lora_config = LoraConfig(
                 r=lora_r,
                 lora_alpha=lora_alpha,
                 target_modules=lora_target_modules,
                 lora_dropout=lora_dropout,
                 bias="none",
-                task_type="SEQ_CLS"
+                task_type="FEATURE_EXTRACTION"
             )
-            print("DEBUG: model class =", type(self.model))
-            for name, module in self.model.named_modules():
-                if "proj" in name:
-                    print("DEBUG: found module", name, "->", module)
             self.model = get_peft_model(self.model, lora_config)
-            print("LoRA enabled! Model type:", type(self.model))
-            for name, param in self.model.named_parameters():
-                if "lora" in name:
-                    print("LoRA param:", name, param.shape)
+            print(f"LoRA enabled: r={lora_r}, alpha={lora_alpha}, target_modules={lora_target_modules}")
 
     def forward(self, words: list[list[str]]) -> Tensor:
         """
@@ -84,8 +77,7 @@ class WordTransformerEncoder(nn.Module):
         ])
 
         # Run model and extract subtokens embeddings from the last layer.
-        outputs = self.model(**subtokens, output_hidden_states=True)
-        subtokens_embeddings = outputs.hidden_states[-1] 
+        subtokens_embeddings = self.model(**subtokens).last_hidden_state
 
         # Aggreate subtokens embeddings into words embeddings.
         # [batch_size, n_words, embedding_size]
@@ -134,7 +126,7 @@ class WordTransformerEncoder(nn.Module):
 
     def get_embeddings_layer(self):
         """Returns the embeddings model."""
-        return self.model.roberta.embeddings
+        return self.model.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:b27ba6b77caece9754d522ddc5bf9e63844a10a04e972b65961a9baccbb08bf5
-size 1134198480
+oid sha256:417921753fb771613766e89226a2fcbdcd259ac9a4c9acbfa55cce7ccb4e1222
+size 1134190536