Upload folder using huggingface_hub

config.json

@@ -19,6 +19,7 @@
   "predict_char": true,
   "predict_length": true,
   "predict_path": true,
+  "predict_path_uncertainty": true,
   "sep_token_id": 2,
   "transformers_version": "4.57.3",
   "unk_token_id": 4,

configuration_swipe.py

@@ -22,6 +22,8 @@ class SwipeTransformerConfig(PretrainedConfig):
         max_char_len (int, optional): Maximum character sequence length. Defaults to 38.
         path_input_dim (int, optional): Path feature dimension. Defaults to 6 for (x, y, dx, dy, ds, log_dt).
         predict_path (bool, optional): Whether to predict path coordinates. Defaults to True.
+        predict_path_uncertainty (bool, optional): Whether to predict log sigma for path coords.
+            Defaults to False.
         pad_token_id (int, optional): Padding token ID. Defaults to 0.
         cls_token_id (int, optional): CLS token ID. Defaults to 1.
         sep_token_id (int, optional): SEP token ID. Defaults to 2.
@@ -45,6 +47,7 @@ class SwipeTransformerConfig(PretrainedConfig):
         path_input_dim: int = 6,
         predict_char: bool = True,
         predict_path: bool = True,
+        predict_path_uncertainty: bool = False,
         predict_length: bool = True,
         pad_token_id: int = 0,
         cls_token_id: int = 1,
@@ -72,6 +75,7 @@ class SwipeTransformerConfig(PretrainedConfig):
         # Model capabilities
         self.predict_char = predict_char
         self.predict_path = predict_path
+        self.predict_path_uncertainty = predict_path_uncertainty
         self.predict_length = predict_length
 
         # Special tokens

heads.py

@@ -58,7 +58,10 @@ class PathPredictionHead(nn.Module):
         # - ds is non-negative
         # - log_dt is non-negative
         if features.shape[-1] == 6:
-            x_y = torch.sigmoid(features[..., 0:2])
+            # Use sigmoid(2x) to avoid center bias that sigmoid(x) has
+            # Mathematical identity: sigmoid(2x) = 0.5(tanh(x)+1)
+            # The 2x scaling provides steeper gradients, helping escape center attraction
+            x_y = torch.sigmoid(2.0 * features[..., 0:2])
             dx_dy = torch.tanh(features[..., 2:4])
             ds = torch.nn.functional.softplus(features[..., 4:5])
             log_dt = torch.nn.functional.softplus(features[..., 5:6])
@@ -68,6 +71,30 @@ class PathPredictionHead(nn.Module):
         return features
 
 
+class PathUncertaintyHead(nn.Module):
+    """Prediction head for log sigma of path coordinates."""
+
+    def __init__(self, d_model: int, output_dim: int = 6):
+        super().__init__()
+        self.dense = nn.Linear(d_model, d_model)
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.decoder = nn.Linear(d_model, output_dim)
+        self.activation = nn.GELU()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [batch, seq_len, d_model]
+
+        Returns:
+            [batch, seq_len, output_dim] log sigma values.
+        """
+        x = self.dense(hidden_states)
+        x = self.activation(x)
+        x = self.layer_norm(x)
+        return self.decoder(x)
+
+
 class LengthPredictionHead(nn.Module):
     """Regress sequence length (e.g., swipable character count) from CLS embedding."""
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:452880eddd71bf06c0c61750ae3f5ba65670c15aa97da8233df856ac74bd3e56
-size 348207344
+oid sha256:ffae9d3acc0266e3856bc25d739afd98000f3b09aa5c1bec6e0c4e9503b6ddbf
+size 350594936

modeling_swipe.py

@@ -22,6 +22,8 @@ class SwipeTransformerOutput(ModelOutput):
             Prediction scores of the character prediction head (text segment only).
         path_logits (`torch.FloatTensor` of shape `(batch_size, path_length, path_input_dim)`, *optional*):
             Prediction scores of the path prediction head (path segment only, if enabled).
+        path_log_sigma (`torch.FloatTensor` of shape `(batch_size, path_length, path_input_dim)`, *optional*):
+            Predicted log sigma for path coordinates (path segment only, if enabled).
         length_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
             Predicted length from the length head (if enabled).
         last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
@@ -39,6 +41,7 @@ class SwipeTransformerOutput(ModelOutput):
     loss: torch.FloatTensor | None = None
     char_logits: torch.FloatTensor | None = None
     path_logits: torch.FloatTensor | None = None
+    path_log_sigma: torch.FloatTensor | None = None
     length_logits: torch.FloatTensor | None = None
     last_hidden_state: torch.FloatTensor | None = None
     pooler_output: torch.FloatTensor | None = None
@@ -88,7 +91,12 @@ class SwipeTransformerModel(SwipeTransformerPreTrainedModel):
 
         # Import existing components
         from .embeddings import MixedEmbedding
-        from .heads import CharacterPredictionHead, LengthPredictionHead, PathPredictionHead
+        from .heads import (
+            CharacterPredictionHead,
+            LengthPredictionHead,
+            PathPredictionHead,
+            PathUncertaintyHead,
+        )
 
         # Embeddings
         self.embeddings = MixedEmbedding(
@@ -130,8 +138,14 @@ class SwipeTransformerModel(SwipeTransformerPreTrainedModel):
             self.path_head = PathPredictionHead(
                 d_model=config.d_model, output_dim=config.path_input_dim
             )
+            self.path_log_sigma_head = (
+                PathUncertaintyHead(d_model=config.d_model, output_dim=config.path_input_dim)
+                if config.predict_path_uncertainty
+                else None
+            )
         else:
             self.path_head = None
+            self.path_log_sigma_head = None
 
         # Length prediction head (predicts word length from path)
         # Max length is max_char_len (including EOS)
@@ -142,87 +156,37 @@ class SwipeTransformerModel(SwipeTransformerPreTrainedModel):
         # Initialize weights
         self.post_init()
 
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        path_coords: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        labels: torch.Tensor | dict | None = None,
-        return_dict: bool | None = None,
-        output_hidden_states: bool | None = None,
-        output_attentions: bool | None = None,
-        **kwargs,
-    ):
-        """
-        Forward pass of the model.
-
-        Args:
-            input_ids (torch.Tensor): Character token IDs [batch, char_len]
-            path_coords (torch.Tensor): Path features [batch, path_len, path_input_dim]
-                                       Default: [batch, path_len, 6] for (x, y, dx, dy, ds, log_dt)
-            attention_mask (torch.Tensor, optional): Attention mask [batch, seq_len]
-            labels (torch.Tensor or dict, optional): Labels for loss calculation
-                Can be tensor [batch, char_len] or dict with keys like char_labels, path_labels
-            return_dict (bool, optional): Whether to return ModelOutput object
-            output_hidden_states (bool, optional): Whether to output hidden states
-            output_attentions (bool, optional): Whether to output attention weights
-            **kwargs: Additional arguments (for compatibility)
-
-        Returns:
-            SwipeTransformerOutput or tuple: Model outputs with:
-                - loss: Optional loss value
-                - char_logits: Character prediction logits [batch, char_len, vocab_size] (if enabled)
-                - path_logits: Path prediction logits [batch, path_len, path_input_dim] (if enabled)
-                - length_logits: Length regression output [batch] (if enabled)
-                - last_hidden_state: Hidden states [batch, seq_len, d_model]
-                - pooler_output: SEP token embedding [batch, d_model] for similarity/embedding tasks
-                - hidden_states: Tuple of per-layer hidden states (if output_hidden_states=True)
-                - attentions: Tuple of per-layer attention weights (if output_attentions=True)
-        """
-        # Validate required inputs
-        if input_ids is None or path_coords is None:
-            raise ValueError("Both input_ids and path_coords are required")
-
-        # Extract labels if dict (used by custom trainers)
-        if isinstance(labels, dict):
-            char_labels = labels.get("char_labels")
-            # Can handle other label types in the future (path_labels, etc.)
-        else:
-            char_labels = labels
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        output_hidden_states = (
-            output_hidden_states
-            if output_hidden_states is not None
-            else self.config.output_hidden_states
-        )
-        output_attentions = (
-            output_attentions if output_attentions is not None else self.config.output_attentions
-        )
-
-        batch_size = path_coords.shape[0]
-        device = path_coords.device
-
-        # Create [CLS] and [SEP] tokens
+    def _make_special_tokens(
+        self, batch_size: int, *, device: torch.device
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         cls_token = torch.full(
-            (batch_size, 1), fill_value=self.config.cls_token_id, dtype=torch.long, device=device
+            (batch_size, 1),
+            fill_value=self.config.cls_token_id,
+            dtype=torch.long,
+            device=device,
         )
         sep_token = torch.full(
-            (batch_size, 1), fill_value=self.config.sep_token_id, dtype=torch.long, device=device
+            (batch_size, 1),
+            fill_value=self.config.sep_token_id,
+            dtype=torch.long,
+            device=device,
         )
+        return cls_token, sep_token
 
-        # Get embeddings
-        embeddings = self.embeddings(path_coords, input_ids, cls_token, sep_token)
-
-        # Prepare attention mask for encoder
-        if attention_mask is not None:
-            # Convert attention mask: 1 = attend, 0 = ignore
-            # PyTorch expects: False = attend, True = ignore
-            src_key_padding_mask = attention_mask == 0
-        else:
-            src_key_padding_mask = None
+    def _src_key_padding_mask(self, attention_mask: torch.Tensor | None) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        return attention_mask == 0
 
-        # Encode while optionally capturing attentions and per-layer hidden states.
+    def _encode(
+        self,
+        embeddings: torch.Tensor,
+        *,
+        src_key_padding_mask: torch.Tensor | None,
+        output_hidden_states: bool,
+        output_attentions: bool,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...] | None, tuple[torch.Tensor, ...] | None]:
+        """Run encoder with optional per-layer hidden-state + attention capture."""
         attentions: tuple[torch.Tensor, ...] | None = None
         hidden_states_by_layer: list[torch.Tensor] | None = [] if output_hidden_states else None
 
@@ -296,69 +260,165 @@ class SwipeTransformerModel(SwipeTransformerPreTrainedModel):
                 if idx in original_forwards:
                     layer.self_attn.forward = original_forwards[idx]
 
-        path_len = path_coords.shape[1]
-        char_len = input_ids.shape[1]
+        all_hidden_states = None
+        if hidden_states_by_layer is not None:
+            all_hidden_states = (embeddings,) + tuple(hidden_states_by_layer)
+
+        return hidden_states, all_hidden_states, attentions
 
-        # Character prediction (text segment only)
+    def _heads(
+        self,
+        hidden_states: torch.Tensor,
+        *,
+        path_len: int,
+        char_len: int,
+    ) -> tuple[
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor | None,
+        torch.Tensor,
+    ]:
         char_logits = None
         if self.char_head is not None:
-            # Sequence is: [CLS] + path + [SEP] + chars
             char_start = 1 + path_len + 1
             char_hidden = hidden_states[:, char_start : char_start + char_len, :]
             char_logits = self.char_head(char_hidden)
 
-        # Path prediction (path segment only, if enabled)
         path_logits = None
+        path_log_sigma = None
         if self.path_head is not None:
             path_hidden = hidden_states[:, 1 : 1 + path_len, :]
             path_logits = self.path_head(path_hidden)
+            if self.path_log_sigma_head is not None:
+                path_log_sigma = self.path_log_sigma_head(path_hidden)
 
-        # Length prediction from CLS token
-        cls_hidden = hidden_states[:, 0, :]  # [batch, d_model] - CLS at position 0
+        cls_hidden = hidden_states[:, 0, :]
         length_logits = self.length_head(cls_hidden) if self.length_head is not None else None
 
-        # Extract SEP token embedding for pooler output (embeddings/similarity tasks)
-        # SEP is at position 1 + path_len
         sep_position = 1 + path_len
-        pooler_output = hidden_states[:, sep_position, :]  # [batch, d_model]
-
-        # Compute loss if labels provided (masked-only; -100 = ignore)
-        loss = None
-        if char_labels is not None and self.char_head is not None:
-            # Predict only the text segment
-            char_pred = char_logits  # [B, char_len, V]
-            labels_flat = char_labels.reshape(-1)
-            mask = labels_flat != -100
-            if mask.any():
-                logits_flat = char_pred.reshape(-1, self.config.vocab_size)[mask]
-                labels_flat = labels_flat[mask]
-                loss = nn.functional.cross_entropy(logits_flat, labels_flat, reduction="mean")
-            else:
-                loss = torch.tensor(0.0, device=hidden_states.device)
+        pooler_output = hidden_states[:, sep_position, :]
+        return char_logits, path_logits, path_log_sigma, length_logits, pooler_output
+
+    def _char_loss(
+        self,
+        *,
+        char_logits: torch.Tensor | None,
+        char_labels: torch.Tensor | None,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if char_labels is None or self.char_head is None or char_logits is None:
+            return None
+
+        labels_flat = char_labels.reshape(-1)
+        mask = labels_flat != -100
+        if mask.any():
+            logits_flat = char_logits.reshape(-1, self.config.vocab_size)[mask]
+            labels_flat = labels_flat[mask]
+            return nn.functional.cross_entropy(logits_flat, labels_flat, reduction="mean")
+        return torch.tensor(0.0, device=device)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        path_coords: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | dict | None = None,
+        return_dict: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        **kwargs,
+    ):
+        """
+        Forward pass of the model.
+
+        Args:
+            input_ids (torch.Tensor): Character token IDs [batch, char_len]
+            path_coords (torch.Tensor): Path features [batch, path_len, path_input_dim]
+                                       Default: [batch, path_len, 6] for (x, y, dx, dy, ds, log_dt)
+            attention_mask (torch.Tensor, optional): Attention mask [batch, seq_len]
+            labels (torch.Tensor or dict, optional): Labels for loss calculation
+                Can be tensor [batch, char_len] or dict with keys like char_labels, path_labels
+            return_dict (bool, optional): Whether to return ModelOutput object
+            output_hidden_states (bool, optional): Whether to output hidden states
+            output_attentions (bool, optional): Whether to output attention weights
+            **kwargs: Additional arguments (for compatibility)
+
+        Returns:
+            SwipeTransformerOutput or tuple: Model outputs with:
+                - loss: Optional loss value
+                - char_logits: Character prediction logits [batch, char_len, vocab_size] (if enabled)
+                - path_logits: Path prediction logits [batch, path_len, path_input_dim] (if enabled)
+                - path_log_sigma: Path log sigma [batch, path_len, path_input_dim] (if enabled)
+                - length_logits: Length regression output [batch] (if enabled)
+                - last_hidden_state: Hidden states [batch, seq_len, d_model]
+                - pooler_output: SEP token embedding [batch, d_model] for similarity/embedding tasks
+                - hidden_states: Tuple of per-layer hidden states (if output_hidden_states=True)
+                - attentions: Tuple of per-layer attention weights (if output_attentions=True)
+        """
+        if input_ids is None or path_coords is None:
+            raise ValueError("Both input_ids and path_coords are required")
+
+        if isinstance(labels, dict):
+            char_labels = labels.get("char_labels")
+        else:
+            char_labels = labels
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+
+        batch_size = int(path_coords.shape[0])
+        device = path_coords.device
+        cls_token, sep_token = self._make_special_tokens(batch_size, device=device)
+        embeddings = self.embeddings(path_coords, input_ids, cls_token, sep_token)
+
+        src_key_padding_mask = self._src_key_padding_mask(attention_mask)
+        hidden_states, all_hidden_states, attentions = self._encode(
+            embeddings,
+            src_key_padding_mask=src_key_padding_mask,
+            output_hidden_states=bool(output_hidden_states),
+            output_attentions=bool(output_attentions),
+        )
+
+        path_len = int(path_coords.shape[1])
+        char_len = int(input_ids.shape[1])
+        char_logits, path_logits, path_log_sigma, length_logits, pooler_output = self._heads(
+            hidden_states,
+            path_len=path_len,
+            char_len=char_len,
+        )
+
+        loss = self._char_loss(
+            char_logits=char_logits,
+            char_labels=char_labels,
+            device=hidden_states.device,
+        )
 
         if not return_dict:
-            hidden_tuple = None
-            if hidden_states_by_layer is not None:
-                hidden_tuple = (embeddings,) + tuple(hidden_states_by_layer)
             output = (
                 char_logits,
                 path_logits,
                 length_logits,
                 hidden_states,
                 pooler_output,
-                hidden_tuple,
+                all_hidden_states,
                 attentions,
+                path_log_sigma,
             )
             return (loss,) + output if loss is not None else output
 
-        all_hidden_states = None
-        if hidden_states_by_layer is not None:
-            all_hidden_states = (embeddings,) + tuple(hidden_states_by_layer)
-
         return SwipeTransformerOutput(
             loss=loss,
             char_logits=char_logits,
             path_logits=path_logits,
+            path_log_sigma=path_log_sigma,
             length_logits=length_logits,
             last_hidden_state=hidden_states,
             pooler_output=pooler_output,

processing_swipe.py

@@ -40,11 +40,7 @@ class SwipeProcessor(ProcessorMixin):
         self.max_char_len = max_char_len
         self.path_input_dim = path_input_dim
         self.path_resample_mode = path_resample_mode
-        # Attributes expected by newer transformers (not used for swipe models)
         self.chat_template = None
-        self.audio_tokenizer = None
-        self.feature_extractor = None
-        self.image_processor = None
 
     def __call__(
         self,
@@ -94,47 +90,87 @@ class SwipeProcessor(ProcessorMixin):
         if path_coords is None and text is None:
             raise ValueError("Must provide either path_coords or text (or both)")
 
-        # Determine batch size
+        batch_size, path_coords, text = self._infer_batch_size(path_coords, text)
+
+        result: dict[str, Any] = {}
+
+        path_coords_out, path_mask = self._process_path_coords(
+            path_coords=path_coords,
+            batch_size=batch_size,
+            truncation=truncation,
+            padding=padding,
+            return_tensors=return_tensors,
+        )
+        result["path_coords"] = path_coords_out
+
+        input_ids, char_mask = self._process_text(
+            text=text,
+            batch_size=batch_size,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            **kwargs,
+        )
+        result["input_ids"] = input_ids
+
+        result["attention_mask"] = self._build_attention_mask(
+            path_mask=path_mask,
+            char_mask=char_mask,
+            batch_size=batch_size,
+            return_tensors=return_tensors,
+        )
+
+        self._convert_result_in_place(result, return_tensors=return_tensors)
+        return result
+
+    def _infer_batch_size(
+        self,
+        path_coords: (
+            list[dict[str, float]]
+            | list[list[dict[str, float]]]
+            | list[list[list[float]]]
+            | torch.Tensor
+            | np.ndarray
+            | None
+        ),
+        text: str | list[str] | None,
+    ) -> tuple[int, Any, str | list[str] | None]:
         if path_coords is not None:
-            # Handle path coordinates
             if isinstance(path_coords, (list, tuple)):
                 if len(path_coords) == 0:
                     batch_size = 1
                 else:
                     first = path_coords[0]
-                    # Raw single path: [{"x","y","t"}, ...]
                     if isinstance(first, dict):
                         batch_size = 1
-                    # Raw batch of paths: [[{"x","y","t"}, ...], ...]
                     elif (
                         isinstance(first, (list, tuple))
                         and len(first) > 0
                         and isinstance(first[0], dict)
                     ):
                         batch_size = len(path_coords)
-                    # Numeric batch: [[[...], ...], ...] where points are lists/tuples
                     elif (
                         isinstance(first, (list, tuple))
                         and len(first) > 0
                         and isinstance(first[0], (list, tuple))
                     ):
                         path_coords = torch.tensor(path_coords, dtype=torch.float32)
-                        batch_size = path_coords.shape[0]
+                        batch_size = int(path_coords.shape[0])
                     else:
-                        # Numeric single path: [[...], [...], ...]
                         path_coords = torch.tensor([path_coords], dtype=torch.float32)
-                        batch_size = path_coords.shape[0]
+                        batch_size = int(path_coords.shape[0])
             elif isinstance(path_coords, np.ndarray):
                 path_coords = torch.from_numpy(path_coords).float()
                 if path_coords.dim() == 2:
-                    # Single path, add batch dimension
                     path_coords = path_coords.unsqueeze(0)
-                batch_size = path_coords.shape[0]
+                batch_size = int(path_coords.shape[0])
             elif isinstance(path_coords, torch.Tensor):
                 if path_coords.dim() == 2:
-                    # Single path, add batch dimension
                     path_coords = path_coords.unsqueeze(0)
-                batch_size = path_coords.shape[0]
+                batch_size = int(path_coords.shape[0])
+            else:
+                batch_size = 1
         elif text is not None:
             if isinstance(text, str):
                 batch_size = 1
@@ -144,238 +180,252 @@ class SwipeProcessor(ProcessorMixin):
         else:
             batch_size = 1
 
-        result = {}
+        return batch_size, path_coords, text
 
-        # Process path coordinates
-        if path_coords is not None:
-            # Check if path_coords is raw data (list of dicts) or already a tensor
-            if isinstance(path_coords, (list, tuple)) and len(path_coords) > 0:
-                first_elem = path_coords[0]
-
-                # Raw single path: [{"x","y","t"}, ...]
-                if isinstance(first_elem, dict) and "x" in first_elem:
+    def _process_path_coords(
+        self,
+        *,
+        path_coords,
+        batch_size: int,
+        truncation: bool,
+        padding: bool | str,
+        return_tensors: str | None,
+    ) -> tuple[Any, Any]:
+        if path_coords is None:
+            path_coords_out = torch.zeros(batch_size, self.max_path_len, self.path_input_dim)
+            path_mask = torch.zeros(batch_size, self.max_path_len, dtype=torch.long)
+            return path_coords_out, path_mask
+
+        if isinstance(path_coords, (list, tuple)) and len(path_coords) > 0:
+            first_elem = path_coords[0]
+
+            if isinstance(first_elem, dict) and "x" in first_elem:
+                path_feats, mask = preprocess_raw_path_to_features(
+                    path_coords,
+                    self.max_path_len,
+                    resample_mode=self.path_resample_mode,
+                )
+                if return_tensors == "pt":
+                    return (
+                        torch.from_numpy(path_feats).float().unsqueeze(0),
+                        torch.from_numpy(mask).long().unsqueeze(0),
+                    )
+                return (np.expand_dims(path_feats, axis=0), np.expand_dims(mask, axis=0))
+
+            if (
+                isinstance(first_elem, (list, tuple))
+                and len(first_elem) > 0
+                and isinstance(first_elem[0], dict)
+                and "x" in first_elem[0]
+            ):
+                processed_paths = []
+                path_masks = []
+                for path in path_coords:
                     path_feats, mask = preprocess_raw_path_to_features(
-                        path_coords,
+                        path,
                         self.max_path_len,
                         resample_mode=self.path_resample_mode,
                     )
-                    if return_tensors == "pt":
-                        path_coords = torch.from_numpy(path_feats).float().unsqueeze(0)
-                        _path_mask = torch.from_numpy(mask).long().unsqueeze(0)
-                    else:
-                        path_coords = np.expand_dims(path_feats, axis=0)
-                        _path_mask = np.expand_dims(mask, axis=0)
-
-                # Raw batch of paths: [[{"x","y","t"}, ...], ...]
-                elif (
-                    isinstance(first_elem, (list, tuple))
-                    and len(first_elem) > 0
-                    and isinstance(first_elem[0], dict)
-                    and "x" in first_elem[0]
-                ):
-                    processed_paths = []
-                    path_masks = []
-                    for path in path_coords:
-                        path_feats, mask = preprocess_raw_path_to_features(
-                            path,
-                            self.max_path_len,
-                            resample_mode=self.path_resample_mode,
-                        )
-                        processed_paths.append(path_feats)
-                        path_masks.append(mask)
-
-                    path_coords = np.stack(processed_paths)  # [batch, max_path_len, 6]
-                    _path_mask = np.stack(path_masks)  # [batch, max_path_len]
-
-                    if return_tensors == "pt":
-                        path_coords = torch.from_numpy(path_coords).float()
-                        _path_mask = torch.from_numpy(_path_mask).long()
-
-                else:
-                    # Numeric list input; process as before
-                    path_coords = torch.tensor(path_coords, dtype=torch.float32)
-                    if path_coords.dim() == 2:
-                        path_coords = path_coords.unsqueeze(0)
-
-                    current_path_len = path_coords.shape[1]
-                    if truncation and current_path_len > self.max_path_len:
-                        path_coords = path_coords[:, : self.max_path_len, :]
-                    if padding and current_path_len < self.max_path_len:
-                        pad_len = self.max_path_len - current_path_len
-                        pad_shape = (batch_size, pad_len, self.path_input_dim)
-                        path_coords = torch.cat([path_coords, torch.zeros(pad_shape)], dim=1)
-
-                    _path_mask = torch.ones(batch_size, self.max_path_len, dtype=torch.long)
-                    is_padding = (path_coords == 0).all(dim=-1)
-                    _path_mask[is_padding] = 0
-            elif isinstance(path_coords, np.ndarray):
-                path_coords = torch.from_numpy(path_coords).float()
-                if path_coords.dim() == 2:
-                    path_coords = path_coords.unsqueeze(0)
-                # If user provided raw (x,y,t) triples but model expects engineered features,
-                # convert to motion features and resample.
-                if path_coords.shape[-1] == 3 and self.path_input_dim == 6:
-                    processed_paths = []
-                    path_masks = []
-                    for path in path_coords.cpu().numpy():
-                        raw = [{"x": float(p[0]), "y": float(p[1]), "t": float(p[2])} for p in path]
-                        path_feats, mask = preprocess_raw_path_to_features(
-                            raw,
-                            self.max_path_len,
-                            resample_mode=self.path_resample_mode,
-                        )
-                        processed_paths.append(path_feats)
-                        path_masks.append(mask)
-
-                    path_coords = torch.from_numpy(np.stack(processed_paths)).float()
-                    _path_mask = torch.from_numpy(np.stack(path_masks)).long()
-                else:
-                    _path_mask = torch.ones(
-                        path_coords.shape[0], self.max_path_len, dtype=torch.long
-                    )
-            elif isinstance(path_coords, torch.Tensor):
-                if path_coords.dim() == 2:
-                    path_coords = path_coords.unsqueeze(0)
-                # If user provided raw (x,y,t) triples but model expects engineered features,
-                # convert to motion features and resample.
-                if path_coords.shape[-1] == 3 and self.path_input_dim == 6:
-                    processed_paths = []
-                    path_masks = []
-                    for path in path_coords.detach().cpu().numpy():
-                        raw = [{"x": float(p[0]), "y": float(p[1]), "t": float(p[2])} for p in path]
-                        path_feats, mask = preprocess_raw_path_to_features(
-                            raw,
-                            self.max_path_len,
-                            resample_mode=self.path_resample_mode,
-                        )
-                        processed_paths.append(path_feats)
-                        path_masks.append(mask)
-
-                    path_coords = torch.from_numpy(np.stack(processed_paths)).float()
-                    _path_mask = torch.from_numpy(np.stack(path_masks)).long()
-                else:
-                    _path_mask = torch.ones(
-                        path_coords.shape[0], self.max_path_len, dtype=torch.long
+                    processed_paths.append(path_feats)
+                    path_masks.append(mask)
+
+                path_coords_np = np.stack(processed_paths)
+                path_mask_np = np.stack(path_masks)
+                if return_tensors == "pt":
+                    return torch.from_numpy(path_coords_np).float(), torch.from_numpy(
+                        path_mask_np
+                    ).long()
+                return path_coords_np, path_mask_np
+
+            # Numeric list input
+            path_tensor = torch.tensor(path_coords, dtype=torch.float32)
+            if path_tensor.dim() == 2:
+                path_tensor = path_tensor.unsqueeze(0)
+
+            current_path_len = int(path_tensor.shape[1])
+            if truncation and current_path_len > self.max_path_len:
+                path_tensor = path_tensor[:, : self.max_path_len, :]
+            if padding and current_path_len < self.max_path_len:
+                pad_len = self.max_path_len - current_path_len
+                pad_shape = (batch_size, pad_len, self.path_input_dim)
+                path_tensor = torch.cat([path_tensor, torch.zeros(pad_shape)], dim=1)
+
+            path_mask = torch.ones(batch_size, self.max_path_len, dtype=torch.long)
+            is_padding = (path_tensor == 0).all(dim=-1)
+            path_mask[is_padding] = 0
+            return path_tensor, path_mask
+
+        if isinstance(path_coords, np.ndarray):
+            path_coords = torch.from_numpy(path_coords).float()
+
+        if isinstance(path_coords, torch.Tensor):
+            if path_coords.dim() == 2:
+                path_coords = path_coords.unsqueeze(0)
+            if path_coords.shape[-1] == 3 and self.path_input_dim == 6:
+                processed_paths = []
+                path_masks = []
+                for path in path_coords.detach().cpu().numpy():
+                    raw = [{"x": float(p[0]), "y": float(p[1]), "t": float(p[2])} for p in path]
+                    path_feats, mask = preprocess_raw_path_to_features(
+                        raw,
+                        self.max_path_len,
+                        resample_mode=self.path_resample_mode,
                     )
+                    processed_paths.append(path_feats)
+                    path_masks.append(mask)
+                return torch.from_numpy(np.stack(processed_paths)).float(), torch.from_numpy(
+                    np.stack(path_masks)
+                ).long()
+
+            if int(path_coords.shape[-1]) != int(self.path_input_dim):
+                raise ValueError(
+                    f"Expected path_coords.shape[-1] == path_input_dim ({self.path_input_dim}), "
+                    f"got {int(path_coords.shape[-1])}. If your path is (x,y,t), pass D=3."
+                )
 
-            result["path_coords"] = path_coords
-        else:
-            # No path coords provided, create empty/zero tensors
-            path_coords = torch.zeros(batch_size, self.max_path_len, self.path_input_dim)
-            _path_mask = torch.zeros(batch_size, self.max_path_len, dtype=torch.long)
-            result["path_coords"] = path_coords
-
-        # Process text
-        if text is not None:
-            # Ensure text is a list
-            if isinstance(text, str):
-                text = [text]
-
-            # Tokenize text
-            text_max_length = max_length if max_length is not None else self.max_char_len
-
-            # First tokenize without padding/truncation to add EOS
-            encoded_raw = self.tokenizer(
-                text,
-                padding=False,
-                truncation=False,
-                return_tensors=None,  # Get lists first
-                **kwargs,
+            path_tensor = path_coords
+            current_path_len = int(path_tensor.shape[1])
+            if truncation and current_path_len > self.max_path_len:
+                path_tensor = path_tensor[:, : self.max_path_len, :]
+            if padding and current_path_len < self.max_path_len:
+                pad_len = self.max_path_len - current_path_len
+                pad_shape = (int(path_tensor.shape[0]), pad_len, int(path_tensor.shape[-1]))
+                pad = torch.zeros(pad_shape, dtype=path_tensor.dtype, device=path_tensor.device)
+                path_tensor = torch.cat([path_tensor, pad], dim=1)
+
+            path_mask = torch.ones(
+                int(path_tensor.shape[0]),
+                int(path_tensor.shape[1]),
+                dtype=torch.long,
+                device=path_tensor.device,
             )
+            is_padding = (path_tensor == 0).all(dim=-1)
+            path_mask[is_padding] = 0
+            return path_tensor, path_mask
 
-            # Add EOS token after each word (matching training dataset behavior)
-            eos_id = self.tokenizer.eos_token_id
-            for i in range(len(encoded_raw["input_ids"])):
-                # Add EOS if not already present
-                if encoded_raw["input_ids"][i][-1] != eos_id:
-                    encoded_raw["input_ids"][i].append(eos_id)
-
-            # Now apply padding and truncation
-            max_len_needed = max(len(ids) for ids in encoded_raw["input_ids"])
-            if truncation and max_len_needed > text_max_length:
-                # Truncate but preserve EOS at the end
-                for i in range(len(encoded_raw["input_ids"])):
-                    if len(encoded_raw["input_ids"][i]) > text_max_length:
-                        encoded_raw["input_ids"][i] = encoded_raw["input_ids"][i][
-                            : text_max_length - 1
-                        ] + [eos_id]
-
-            # Pad sequences
-            if padding:
-                pad_id = self.tokenizer.pad_token_id
-                for i in range(len(encoded_raw["input_ids"])):
-                    seq_len = len(encoded_raw["input_ids"][i])
-                    if seq_len < text_max_length:
-                        encoded_raw["input_ids"][i].extend([pad_id] * (text_max_length - seq_len))
-
-            # Create attention mask (1 for real tokens + EOS, 0 for padding)
-            _char_mask = []
-            for ids in encoded_raw["input_ids"]:
-                mask = [1 if token_id != self.tokenizer.pad_token_id else 0 for token_id in ids]
-                _char_mask.append(mask)
-
-            # Convert to tensors if requested
-            if return_tensors == "pt":
-                result["input_ids"] = torch.tensor(encoded_raw["input_ids"], dtype=torch.long)
-                _char_mask = torch.tensor(_char_mask, dtype=torch.long)
-            elif return_tensors == "np":
-                result["input_ids"] = np.array(encoded_raw["input_ids"], dtype=np.int64)
-                _char_mask = np.array(_char_mask, dtype=np.int64)
-            else:
-                result["input_ids"] = encoded_raw["input_ids"]
-        else:
-            # No text provided, create padding tokens
+        # Fallback: treat unknown input as empty path.
+        path_coords_out = torch.zeros(batch_size, self.max_path_len, self.path_input_dim)
+        path_mask = torch.zeros(batch_size, self.max_path_len, dtype=torch.long)
+        return path_coords_out, path_mask
+
+    def _process_text(
+        self,
+        *,
+        text: str | list[str] | None,
+        batch_size: int,
+        padding: bool | str,
+        truncation: bool,
+        max_length: int | None,
+        return_tensors: str | None,
+        **kwargs: Any,
+    ) -> tuple[Any, Any]:
+        if text is None:
             if return_tensors == "pt":
                 char_tokens = torch.full(
-                    (batch_size, self.max_char_len), self.tokenizer.pad_token_id, dtype=torch.long
+                    (batch_size, self.max_char_len),
+                    self.tokenizer.pad_token_id,
+                    dtype=torch.long,
                 )
-                _char_mask = torch.zeros(batch_size, self.max_char_len, dtype=torch.long)
+                char_mask = torch.zeros(batch_size, self.max_char_len, dtype=torch.long)
             elif return_tensors == "np":
                 char_tokens = np.full(
-                    (batch_size, self.max_char_len), self.tokenizer.pad_token_id, dtype=np.int64
+                    (batch_size, self.max_char_len),
+                    self.tokenizer.pad_token_id,
+                    dtype=np.int64,
                 )
-                _char_mask = np.zeros((batch_size, self.max_char_len), dtype=np.int64)
+                char_mask = np.zeros((batch_size, self.max_char_len), dtype=np.int64)
             else:
                 char_tokens = [
                     [self.tokenizer.pad_token_id] * self.max_char_len for _ in range(batch_size)
                 ]
-                _char_mask = [[0] * self.max_char_len for _ in range(batch_size)]
+                char_mask = [[0] * self.max_char_len for _ in range(batch_size)]
+            return char_tokens, char_mask
+
+        if isinstance(text, str):
+            text = [text]
+
+        text_max_length = max_length if max_length is not None else self.max_char_len
+
+        encoded_raw = self.tokenizer(
+            text,
+            padding=False,
+            truncation=False,
+            return_tensors=None,
+            **kwargs,
+        )
 
-            result["input_ids"] = char_tokens
+        eos_id = self.tokenizer.eos_token_id
+        for i in range(len(encoded_raw["input_ids"])):
+            if encoded_raw["input_ids"][i][-1] != eos_id:
+                encoded_raw["input_ids"][i].append(eos_id)
+
+        max_len_needed = max(len(ids) for ids in encoded_raw["input_ids"])
+        if truncation and max_len_needed > text_max_length:
+            for i in range(len(encoded_raw["input_ids"])):
+                if len(encoded_raw["input_ids"][i]) > text_max_length:
+                    encoded_raw["input_ids"][i] = encoded_raw["input_ids"][i][
+                        : text_max_length - 1
+                    ] + [eos_id]
+
+        if padding:
+            pad_id = self.tokenizer.pad_token_id
+            for i in range(len(encoded_raw["input_ids"])):
+                seq_len = len(encoded_raw["input_ids"][i])
+                if seq_len < text_max_length:
+                    encoded_raw["input_ids"][i].extend([pad_id] * (text_max_length - seq_len))
+
+        char_mask_list = [
+            [1 if token_id != self.tokenizer.pad_token_id else 0 for token_id in ids]
+            for ids in encoded_raw["input_ids"]
+        ]
 
-        # Create combined attention mask: [CLS] + path + [SEP] + chars
-        # Sequence structure: [CLS:1] + _path_mask + [SEP:1] + _char_mask
+        if return_tensors == "pt":
+            return (
+                torch.tensor(encoded_raw["input_ids"], dtype=torch.long),
+                torch.tensor(char_mask_list, dtype=torch.long),
+            )
+        if return_tensors == "np":
+            return (
+                np.array(encoded_raw["input_ids"], dtype=np.int64),
+                np.array(char_mask_list, dtype=np.int64),
+            )
+        return encoded_raw["input_ids"], char_mask_list
+
+    def _build_attention_mask(
+        self,
+        *,
+        path_mask,
+        char_mask,
+        batch_size: int,
+        return_tensors: str | None,
+    ):
         if return_tensors == "pt":
             cls_mask = torch.ones(batch_size, 1, dtype=torch.long)
             sep_mask = torch.ones(batch_size, 1, dtype=torch.long)
-            attention_mask = torch.cat([cls_mask, _path_mask, sep_mask, _char_mask], dim=1)
-        elif return_tensors == "np":
+            return torch.cat([cls_mask, path_mask, sep_mask, char_mask], dim=1)
+        if return_tensors == "np":
             cls_mask = np.ones((batch_size, 1), dtype=np.int64)
             sep_mask = np.ones((batch_size, 1), dtype=np.int64)
-            attention_mask = np.concatenate([cls_mask, _path_mask, sep_mask, _char_mask], axis=1)
-        else:
-            cls_mask = [[1] for _ in range(batch_size)]
-            sep_mask = [[1] for _ in range(batch_size)]
-            attention_mask = [
-                cls + path.tolist() + sep + char
-                for cls, path, sep, char in zip(
-                    cls_mask, _path_mask, sep_mask, _char_mask, strict=False
-                )
-            ]
-
-        result["attention_mask"] = attention_mask
-
-        # Convert to requested format
+            return np.concatenate([cls_mask, path_mask, sep_mask, char_mask], axis=1)
+
+        cls_mask = [[1] for _ in range(batch_size)]
+        sep_mask = [[1] for _ in range(batch_size)]
+        return [
+            cls + path.tolist() + sep + char
+            for cls, path, sep, char in zip(cls_mask, path_mask, sep_mask, char_mask, strict=False)
+        ]
+
+    def _convert_result_in_place(
+        self, result: dict[str, Any], *, return_tensors: str | None
+    ) -> None:
         if return_tensors == "np":
-            for key in result:
-                if isinstance(result[key], torch.Tensor):
-                    result[key] = result[key].numpy()
+            for key, value in list(result.items()):
+                if isinstance(value, torch.Tensor):
+                    result[key] = value.numpy()
         elif return_tensors is None:
-            for key in result:
-                if isinstance(result[key], torch.Tensor):
-                    result[key] = result[key].tolist()
-
-        return result
+            for key, value in list(result.items()):
+                if isinstance(value, torch.Tensor):
+                    result[key] = value.tolist()
 
     def batch_decode(self, token_ids, **kwargs):
         """

special_tokens_map.json

@@ -1,44 +1,8 @@
 {
-  "cls_token": {
-    "content": "[CLS]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "[EOS]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "mask_token": {
-    "content": "[MASK]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": {
-    "content": "[PAD]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "sep_token": {
-    "content": "[SEP]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "unk_token": {
-    "content": "[UNK]",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  }
+  "cls_token": "[CLS]",
+  "eos_token": "[EOS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
 }

tokenizer_config.json

@@ -49,18 +49,22 @@
       "special": true
     }
   },
+  "additional_special_tokens": null,
   "auto_map": {
     "AutoTokenizer": [
       "tokenization_swipe.SwipeTokenizer",
       null
     ]
   },
+  "backend": "custom",
   "clean_up_tokenization_spaces": false,
   "cls_token": "[CLS]",
   "eos_token": "[EOS]",
-  "extra_special_tokens": {},
+  "extra_special_tokens": [],
+  "is_local": true,
   "mask_token": "[MASK]",
   "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {},
   "pad_token": "[PAD]",
   "processor_class": "SwipeProcessor",
   "sep_token": "[SEP]",