fix bad readme and remove warnings about uninitialized tensors

README.md

@@ -7,20 +7,20 @@ tags:
 library_name: transformers
 paper: https://arxiv.org/abs/2201.05601
 ---
-```python
-from transformers import AutoModel, AutoTokenizer
-
-model = AutoModel.from_pretrained("mideind/IceBERT-PoS", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("mideind/IceBERT-PoS")
-
 ## Prediction Methods
 
 The model provides two prediction methods:
 
-- **`predict_labels_from_text()`**: Returns structured predictions as (category, [attributes]) tuples. Use this for downstream NLP tasks or when you need the semantic meaning of each prediction.
+- **`predict_labels_from_text()`**: Returns structured predictions as (category, [attributes]) tuples. These can be slightly more readable and more suitable for some applications.
 - **`predict_ifd_labels_from_text()`**: Returns predictions in IFD (Icelandic Frequency Dictionary) format. Use this for evaluation against MIM-GOLD datasets or when you need compatibility with traditional Icelandic POS taggers.
 
-# Example sentence which is already tokenized (with a classic tokenizer)
+```python
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained("mideind/IceBERT-PoS", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("mideind/IceBERT-PoS")
+
+# Example sentence
 sentence = "Ég veit að þú kemur í kvöld til mín ."
 
 # Get predictions in (category, [attributes]) format

modeling.py

@@ -91,22 +91,26 @@ class IceBertPosForTokenClassification(PreTrainedModel):
         """Setup label mappings using schema methods."""
         schema = self.config.label_schema
 
-        # Get model device for tensor creation
-        device = next(self.parameters()).device if len(list(self.parameters())) > 0 else torch.device("cpu")
-
-        # Register group mask as buffer so it moves with model.to(device)
-        self.register_buffer("group_mask", schema.get_group_masks(device=device))
-
-        # Register group attribute indices as buffers
-        group_attr_indices = schema.get_group_name_to_group_attr_indices(device=device)
-        self.group_name_to_group_attr_indices = {}
-        for group_name, indices in group_attr_indices.items():
-            buffer_name = f"group_attr_indices_{group_name}"
-            self.register_buffer(buffer_name, indices)
-            self.group_name_to_group_attr_indices[group_name] = getattr(self, buffer_name)
-
+        # Create tensors as regular attributes (not buffers to avoid init warnings)
+        self.group_mask = schema.get_group_masks()
+        self.group_name_to_group_attr_indices = schema.get_group_name_to_group_attr_indices()
+        
         # Category name to index mapping (regular dict, no device movement needed)
         self.category_name_to_index = schema.get_category_name_to_index()
+    
+    def _apply(self, fn):
+        """Override _apply to move our custom tensors with the model."""
+        super()._apply(fn)
+        
+        # Move our custom tensors when model.to(device) is called
+        if hasattr(self, 'group_mask'):
+            self.group_mask = fn(self.group_mask)
+        
+        if hasattr(self, 'group_name_to_group_attr_indices'):
+            for group_name, tensor in self.group_name_to_group_attr_indices.items():
+                self.group_name_to_group_attr_indices[group_name] = fn(tensor)
+        
+        return self
 
     def forward(
         self,