Upload folder using huggingface_hub

README.md

@@ -1,3 +1,155 @@
----
-license: apache-2.0
----
+# SwipeALot Base Model
+
+Multimodal transformer for swipe keyboard prediction. Trained on the [futo-org/swipe.futo.org](https://huggingface.co/datasets/futo-org/swipe.futo.org) dataset.
+
+## Quick Start
+
+```python
+from transformers import AutoModel, AutoProcessor
+import torch
+
+# Load model
+model = AutoModel.from_pretrained("path/to/model", trust_remote_code=True)
+processor = AutoProcessor.from_pretrained("path/to/model", trust_remote_code=True)
+model.eval()
+
+# Example: Predict word from swipe path
+from datasets import load_dataset
+from swipealot.data.dataset import normalize_coordinates, sample_path_points
+
+# Load sample
+dataset = load_dataset("futo-org/swipe.futo.org", split="test[:1]")
+item = dataset[0]
+
+# Preprocess path
+normalized = normalize_coordinates(item["data"], item["canvas_width"], item["canvas_height"])
+path_coords, _ = sample_path_points(normalized, processor.max_path_len)
+path = torch.tensor([path_coords], dtype=torch.float32)
+
+# Get predictions
+inputs = processor(path_coords=path, text=None, return_tensors="pt")
+
+with torch.no_grad():
+    outputs = model(**inputs)
+
+# Length prediction
+predicted_length = outputs.length_logits.argmax(dim=-1).item()
+print(f"Predicted word length: {predicted_length}")
+```
+
+## Model Details
+
+- **Architecture**: Transformer encoder (768-dim, 12 layers, 12 heads)
+- **Parameters**: 87M
+- **Training Data**: futo-org/swipe.futo.org dataset
+- **Max Path Length**: 128 points
+- **Max Word Length**: 48 characters
+- **Vocab Size**: 43 (a-z, 0-9, special tokens)
+
+## Capabilities
+
+### 1. Character Prediction
+Predict characters from swipe paths with partial text context.
+
+**Use Case**: Autocorrection, suggestion ranking
+
+### 2. Length Prediction
+Predict word length from swipe path alone.
+
+**Accuracy**: 89% exact, 96% within ±1
+
+**Use Case**: Pre-filtering candidate words
+
+### 3. Path Reconstruction
+Reconstruct missing path coordinates.
+
+**MSE**: 0.005 on masked points
+
+**Use Case**: Noise reduction, gesture smoothing
+
+### 4. Embedding Extraction
+Extract fixed-size embeddings for similarity search.
+
+**Dimension**: 768
+
+**Use Case**: Similar gesture search, deduplication
+
+## Usage Examples
+
+See the [full documentation](https://github.com/dleemiller/legendary-waffle) for detailed examples of each capability.
+
+### Masked Character Prediction
+
+```python
+# Process full word, then manually mask positions
+inputs = processor(path_coords=path, text="hello", return_tensors="pt")
+mask_token_id = processor.tokenizer.mask_token_id
+char_ids = inputs["input_ids"][0].tolist()
+char_ids[2] = mask_token_id  # Mask 'l' at position 2
+inputs["input_ids"] = torch.tensor([char_ids], dtype=torch.long)
+
+# Model predicts masked character from path + context
+```
+
+### Full Word Reconstruction
+
+```python
+# Process word, then mask all character positions
+inputs = processor(path_coords=path, text="hello", return_tensors="pt")
+char_ids = inputs["input_ids"][0].tolist()
+mask_token_id = processor.tokenizer.mask_token_id
+masked_ids = [mask_token_id if cid != 0 else 0 for cid in char_ids]
+inputs["input_ids"] = torch.tensor([masked_ids], dtype=torch.long)
+
+# Predict from path only - achieves 94% character accuracy
+```
+
+### Length Prediction
+
+```python
+inputs = processor(path_coords=path, text=None, return_tensors="pt")
+predicted_length = outputs.length_logits.argmax(dim=-1).item()
+```
+
+## Performance Metrics
+
+Evaluated on 200 test samples:
+
+| Task | Metric | Score |
+|------|--------|-------|
+| Masked Prediction (30%) | Character Accuracy | 98.7% |
+| | Top-3 Accuracy | 100% |
+| | Word Accuracy | 97.3% |
+| Full Reconstruction (100%) | Character Accuracy | 94% |
+| | Word Accuracy | 83.7% |
+| Length Prediction | Exact Accuracy | 89% |
+| | Within ±1 | 96% |
+| | Within ±2 | 99% |
+| Path Reconstruction | MSE (masked) | 0.005 |
+
+## Model Outputs
+
+```python
+outputs = model(**inputs)
+
+# Available outputs:
+outputs.char_logits       # [batch, seq_len, vocab_size] - Character predictions
+outputs.length_logits     # [batch, max_length] - Length predictions
+outputs.path_logits       # [batch, seq_len, 3] - Path coordinate predictions
+outputs.pooler_output     # [batch, d_model] - SEP token embeddings for similarity
+outputs.last_hidden_state # [batch, seq_len, d_model] - Hidden representations
+```
+
+## Citation
+
+```bibtex
+@software{swipealot2024,
+  title={SwipeALot: Multimodal Swipe Keyboard Transformer},
+  year={2024},
+  url={https://github.com/dleemiller/legendary-waffle}
+}
+```
+
+## License
+
+MIT License

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "SwipeTransformerModel"
+  ],
+  "cls_token_id": 1,
+  "d_ff": 3072,
+  "d_model": 768,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": 5,
+  "mask_token_id": 3,
+  "max_char_len": 48,
+  "max_path_len": 128,
+  "model_type": "swipe_transformer",
+  "n_heads": 12,
+  "n_layers": 12,
+  "pad_token_id": 0,
+  "predict_path": true,
+  "sep_token_id": 2,
+  "transformers_version": "4.57.3",
+  "unk_token_id": 4,
+  "vocab_size": 43,
+  "auto_map": {
+    "AutoConfig": "configuration_swipe.SwipeTransformerConfig",
+    "AutoModel": "modeling_swipe.SwipeTransformerModel"
+  }
+}

configuration_swipe.py

@@ -0,0 +1,148 @@
+"""Configuration classes for SwipeTransformer HuggingFace models."""
+
+from transformers import PretrainedConfig
+
+
+class SwipeTransformerConfig(PretrainedConfig):
+    """
+    Configuration class for SwipeTransformerModel.
+
+    This configuration stores all the parameters needed to instantiate a
+    SwipeTransformerModel. This is the base configuration for the multimodal
+    swipe keyboard transformer that processes path coordinates and text.
+
+    Args:
+        d_model (int, optional): Hidden dimension size. Defaults to 256.
+        n_layers (int, optional): Number of transformer layers. Defaults to 4.
+        n_heads (int, optional): Number of attention heads. Defaults to 4.
+        d_ff (int, optional): Feedforward dimension. Defaults to 1024.
+        dropout (float, optional): Dropout rate. Defaults to 0.1.
+        vocab_size (int, optional): Size of vocabulary. Defaults to 100.
+        max_path_len (int, optional): Maximum path sequence length. Defaults to 64.
+        max_char_len (int, optional): Maximum character sequence length. Defaults to 38.
+        predict_path (bool, optional): Whether to predict path coordinates. Defaults to True.
+        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.
+        mask_token_id (int, optional): MASK token ID. Defaults to 3.
+        unk_token_id (int, optional): Unknown token ID. Defaults to 4.
+        eos_token_id (int, optional): End-of-sequence token ID. Defaults to 5.
+    """
+
+    model_type = "swipe_transformer"
+
+    def __init__(
+        self,
+        d_model: int = 256,
+        n_layers: int = 4,
+        n_heads: int = 4,
+        d_ff: int = 1024,
+        dropout: float = 0.1,
+        vocab_size: int = 100,
+        max_path_len: int = 64,
+        max_char_len: int = 38,
+        predict_path: bool = True,
+        pad_token_id: int = 0,
+        cls_token_id: int = 1,
+        sep_token_id: int = 2,
+        mask_token_id: int = 3,
+        unk_token_id: int = 4,
+        eos_token_id: int = 5,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, eos_token_id=eos_token_id, **kwargs)
+
+        # Model architecture parameters
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.d_ff = d_ff
+        self.dropout = dropout
+
+        # Vocabulary and sequence length
+        self.vocab_size = vocab_size
+        self.max_path_len = max_path_len
+        self.max_char_len = max_char_len
+
+        # Model capabilities
+        self.predict_path = predict_path
+
+        # Special tokens
+        self.cls_token_id = cls_token_id
+        self.sep_token_id = sep_token_id
+        self.mask_token_id = mask_token_id
+        self.unk_token_id = unk_token_id
+
+
+class SwipeCrossEncoderConfig(PretrainedConfig):
+    """
+    Configuration class for SwipeCrossEncoderForSequenceClassification.
+
+    This configuration extends the base SwipeTransformer config for use in
+    cross-encoder tasks (e.g., path-word similarity scoring).
+
+    Args:
+        d_model (int, optional): Hidden dimension size. Defaults to 256.
+        n_layers (int, optional): Number of transformer layers. Defaults to 4.
+        n_heads (int, optional): Number of attention heads. Defaults to 4.
+        d_ff (int, optional): Feedforward dimension. Defaults to 1024.
+        dropout (float, optional): Dropout rate. Defaults to 0.1.
+        vocab_size (int, optional): Size of vocabulary. Defaults to 100.
+        max_path_len (int, optional): Maximum path sequence length. Defaults to 64.
+        max_char_len (int, optional): Maximum character sequence length. Defaults to 38.
+        num_labels (int, optional): Number of classification labels. Defaults to 1.
+        problem_type (str, optional): Problem type ('regression' or 'single_label_classification'). Defaults to "regression".
+        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.
+        mask_token_id (int, optional): MASK token ID. Defaults to 3.
+        unk_token_id (int, optional): Unknown token ID. Defaults to 4.
+        eos_token_id (int, optional): End-of-sequence token ID. Defaults to 5.
+    """
+
+    model_type = "swipe_cross_encoder"
+
+    def __init__(
+        self,
+        d_model: int = 256,
+        n_layers: int = 4,
+        n_heads: int = 4,
+        d_ff: int = 1024,
+        dropout: float = 0.1,
+        vocab_size: int = 100,
+        max_path_len: int = 64,
+        max_char_len: int = 38,
+        num_labels: int = 1,
+        problem_type: str = "regression",
+        pad_token_id: int = 0,
+        cls_token_id: int = 1,
+        sep_token_id: int = 2,
+        mask_token_id: int = 3,
+        unk_token_id: int = 4,
+        eos_token_id: int = 5,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id, num_labels=num_labels, eos_token_id=eos_token_id, **kwargs
+        )
+
+        # Model architecture parameters
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.d_ff = d_ff
+        self.dropout = dropout
+
+        # Vocabulary and sequence length
+        self.vocab_size = vocab_size
+        self.max_path_len = max_path_len
+        self.max_char_len = max_char_len
+
+        # Classification parameters
+        self.problem_type = problem_type
+
+        # Special tokens
+        self.cls_token_id = cls_token_id
+        self.sep_token_id = sep_token_id
+        self.mask_token_id = mask_token_id
+        self.unk_token_id = unk_token_id

conversion_metadata.json

@@ -0,0 +1,16 @@
+{
+  "original_checkpoint": "checkpoints/base_20251213_164813/best.pt",
+  "original_config": "embedded_in_checkpoint",
+  "converted_at": "2025-12-15 08:03:26.074041",
+  "model_type": "base",
+  "vocab_size": 43,
+  "epoch": 38,
+  "global_step": 70000,
+  "metrics": {
+    "loss": 0.03484317846596241,
+    "char_accuracy": 0.9536226987838745,
+    "word_accuracy": 0.771,
+    "char_loss_mean": 0.03484317846596241,
+    "total_loss_mean": 0.03484317846596241
+  }
+}

embeddings.py

@@ -0,0 +1,179 @@
+"""Embedding layers for SwipeTransformer."""
+
+import torch
+import torch.nn as nn
+
+
+class PathEmbedding(nn.Module):
+    """Embeds path coordinates (x, y, t) to d_model dimension."""
+
+    def __init__(self, d_model: int = 256):
+        super().__init__()
+        self.projection = nn.Linear(3, d_model)
+
+    def forward(self, path_coords: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            path_coords: [batch, seq_len, 3] - (x, y, t) coordinates
+
+        Returns:
+            [batch, seq_len, d_model] embeddings
+        """
+        return self.projection(path_coords)
+
+
+class CharacterEmbedding(nn.Module):
+    """Embeds character tokens."""
+
+    def __init__(self, vocab_size: int, d_model: int = 256, padding_idx: int = 0):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=padding_idx)
+
+    def forward(self, char_tokens: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            char_tokens: [batch, seq_len] character token IDs
+
+        Returns:
+            [batch, seq_len, d_model] embeddings
+        """
+        return self.embedding(char_tokens)
+
+
+class PositionalEmbedding(nn.Module):
+    """Learned positional embeddings."""
+
+    def __init__(self, max_position: int, d_model: int = 256):
+        super().__init__()
+        self.embedding = nn.Embedding(max_position, d_model)
+
+    def forward(self, positions: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            positions: [batch, seq_len] position indices
+
+        Returns:
+            [batch, seq_len, d_model] positional embeddings
+        """
+        return self.embedding(positions)
+
+
+class TypeEmbedding(nn.Module):
+    """Token type embeddings to distinguish PATH (0) vs TEXT (1) tokens."""
+
+    def __init__(self, d_model: int = 256):
+        super().__init__()
+        # 0 = PATH, 1 = TEXT
+        self.embedding = nn.Embedding(2, d_model)
+
+    def forward(self, token_types: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            token_types: [batch, seq_len] type indices (0 or 1)
+
+        Returns:
+            [batch, seq_len, d_model] type embeddings
+        """
+        return self.embedding(token_types)
+
+
+class MixedEmbedding(nn.Module):
+    """
+    Combines path and character embeddings with positional and type information.
+    Constructs sequence: [CLS] + path_tokens + [SEP] + char_tokens
+    """
+
+    def __init__(
+        self,
+        vocab_size: int,
+        max_path_len: int,
+        max_char_len: int,
+        d_model: int = 256,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.d_model = d_model
+
+        # Content embeddings
+        self.path_embedding = PathEmbedding(d_model)
+        self.char_embedding = CharacterEmbedding(vocab_size, d_model, padding_idx=0)
+
+        # Positional embeddings
+        max_seq_len = 1 + max_path_len + 1 + max_char_len  # [CLS] + path + [SEP] + chars
+        self.positional_embedding = PositionalEmbedding(max_seq_len, d_model)
+
+        # Type embeddings
+        self.type_embedding = TypeEmbedding(d_model)
+
+        # Layer norm and dropout
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        path_coords: torch.Tensor,
+        char_tokens: torch.Tensor,
+        cls_token: torch.Tensor,
+        sep_token: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Create mixed sequence with embeddings.
+
+        Args:
+            path_coords: [batch, path_len, 3] path coordinates
+            char_tokens: [batch, char_len] character token IDs
+            cls_token: [batch, 1] CLS token IDs
+            sep_token: [batch, 1] SEP token IDs
+
+        Returns:
+            [batch, total_seq_len, d_model] embeddings where
+            total_seq_len = 1 + path_len + 1 + char_len
+        """
+        batch_size = path_coords.shape[0]
+        path_len = path_coords.shape[1]
+        char_len = char_tokens.shape[1]
+        device = path_coords.device
+
+        # Embed [CLS]
+        cls_emb = self.char_embedding(cls_token)  # [batch, 1, d_model]
+
+        # Embed path
+        path_emb = self.path_embedding(path_coords)  # [batch, path_len, d_model]
+
+        # Embed [SEP]
+        sep_emb = self.char_embedding(sep_token)  # [batch, 1, d_model]
+
+        # Embed characters
+        char_emb = self.char_embedding(char_tokens)  # [batch, char_len, d_model]
+
+        # Concatenate: [CLS] + PATH + [SEP] + CHARS
+        sequence = torch.cat(
+            [cls_emb, path_emb, sep_emb, char_emb], dim=1
+        )  # [batch, seq_len, d_model]
+        seq_len = sequence.shape[1]
+
+        # Add positional embeddings
+        positions = (
+            torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
+        )  # [batch, seq_len]
+        pos_emb = self.positional_embedding(positions)
+
+        # Add type embeddings
+        # Type 0 for [CLS] + path + [SEP], Type 1 for chars
+        type_ids = torch.cat(
+            [
+                torch.zeros(
+                    batch_size, 1 + path_len + 1, dtype=torch.long, device=device
+                ),  # [CLS], path, [SEP]
+                torch.ones(batch_size, char_len, dtype=torch.long, device=device),  # chars
+            ],
+            dim=1,
+        )  # [batch, seq_len]
+        type_emb = self.type_embedding(type_ids)
+
+        # Combine: content + position + type
+        embeddings = sequence + pos_emb + type_emb
+        embeddings = self.layer_norm(embeddings)
+        embeddings = self.dropout(embeddings)
+
+        return embeddings

heads.py

@@ -0,0 +1,109 @@
+"""Prediction heads for SwipeTransformer."""
+
+import torch
+import torch.nn as nn
+
+
+class CharacterPredictionHead(nn.Module):
+    """Prediction head for masked characters."""
+
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        self.dense = nn.Linear(d_model, d_model)
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.decoder = nn.Linear(d_model, vocab_size)
+        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, vocab_size] logits
+        """
+        x = self.dense(hidden_states)
+        x = self.activation(x)
+        x = self.layer_norm(x)
+        logits = self.decoder(x)
+        return logits
+
+
+class PathPredictionHead(nn.Module):
+    """Prediction head for masked path coordinates."""
+
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.dense = nn.Linear(d_model, d_model)
+        self.layer_norm = nn.LayerNorm(d_model)
+        self.decoder = nn.Linear(d_model, 3)  # Predict (x, y, t)
+        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, 3] coordinates in [0, 1] range
+        """
+        x = self.dense(hidden_states)
+        x = self.activation(x)
+        x = self.layer_norm(x)
+        coords = self.decoder(x)
+        coords = torch.sigmoid(coords)  # Ensure [0, 1] range
+        return coords
+
+
+class ClassificationHead(nn.Module):
+    """
+    Classification head for cross-encoder.
+
+    Follows SBERT architecture: Dense → GELU → LayerNorm → Linear(→1)
+    Outputs a single similarity score per input.
+    """
+
+    def __init__(self, d_model: int, num_labels: int = 1):
+        super().__init__()
+        self.dense = nn.Linear(d_model, d_model)
+        self.activation = nn.GELU()
+        self.norm = nn.LayerNorm(d_model)
+        self.classifier = nn.Linear(d_model, num_labels)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            features: [batch, d_model] - typically SEP token embeddings
+
+        Returns:
+            [batch, num_labels] similarity scores
+        """
+        x = self.dense(features)
+        x = self.activation(x)  # GELU
+        x = self.norm(x)  # LayerNorm
+        logits = self.classifier(x)  # [batch, 1] or [batch, num_labels]
+        return logits
+
+
+class LengthPredictionHead(nn.Module):
+    """Predict sequence length (e.g., swipable character count) from CLS embedding."""
+
+    def __init__(self, d_model: int, max_length: int):
+        super().__init__()
+        self.dense = nn.Linear(d_model, d_model)
+        self.activation = nn.GELU()
+        self.norm = nn.LayerNorm(d_model)
+        self.classifier = nn.Linear(d_model, max_length + 1)  # classes: 0..max_length
+
+    def forward(self, cls_features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            cls_features: [batch, d_model] CLS embeddings
+
+        Returns:
+            [batch, max_length+1] logits over lengths
+        """
+        x = self.dense(cls_features)
+        x = self.activation(x)
+        x = self.norm(x)
+        return self.classifier(x)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b00a7a35dc99485501db127c4433afee95bd4f0c34d6c6db8ed6c69eec43404b
+size 348336548

modeling_swipe.py

@@ -0,0 +1,516 @@
+"""HuggingFace-compatible model classes for SwipeTransformer."""
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
+    ModelOutput,
+    SequenceClassifierOutput,
+)
+
+from .configuration_swipe import SwipeCrossEncoderConfig, SwipeTransformerConfig
+
+
+@dataclass
+class SwipeTransformerOutput(ModelOutput):
+    """
+    Output type for SwipeTransformerModel.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (character prediction).
+        char_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, vocab_size)`):
+            Prediction scores of the character prediction head.
+        path_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 3)`, *optional*):
+            Prediction scores of the path prediction head (if enabled).
+        length_logits (`torch.FloatTensor` of shape `(batch_size, max_length+1)`, *optional*):
+            Prediction scores of the length prediction head (if enabled).
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
+            SEP token embeddings for similarity/embedding tasks.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    char_logits: torch.FloatTensor = None
+    path_logits: Optional[torch.FloatTensor] = None
+    length_logits: Optional[torch.FloatTensor] = None
+    last_hidden_state: torch.FloatTensor = None
+    pooler_output: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+class SwipeTransformerPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface
+    for downloading and loading pretrained models.
+    """
+
+    config_class = SwipeTransformerConfig
+    base_model_prefix = "swipe_transformer"
+    supports_gradient_checkpointing = False
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            nn.init.xavier_uniform_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+
+
+class SwipeTransformerModel(SwipeTransformerPreTrainedModel):
+    """
+    HuggingFace-compatible SwipeTransformerModel.
+
+    This model reuses the existing components from src/swipealot/models/
+    and wraps them in a HuggingFace-compatible interface.
+
+    Args:
+        config (SwipeTransformerConfig): Model configuration
+    """
+
+    def __init__(self, config: SwipeTransformerConfig):
+        super().__init__(config)
+        self.config = config
+
+        # Import existing components
+        from .embeddings import MixedEmbedding
+        from .heads import CharacterPredictionHead, LengthPredictionHead, PathPredictionHead
+
+        # Embeddings
+        self.embeddings = MixedEmbedding(
+            vocab_size=config.vocab_size,
+            max_path_len=config.max_path_len,
+            max_char_len=config.max_char_len,
+            d_model=config.d_model,
+            dropout=config.dropout,
+        )
+
+        # Transformer encoder
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.n_heads,
+            dim_feedforward=config.d_ff,
+            dropout=config.dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,  # Pre-LayerNorm
+        )
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=config.n_layers,
+            enable_nested_tensor=False,
+        )
+
+        # Prediction heads
+        self.char_head = CharacterPredictionHead(
+            d_model=config.d_model,
+            vocab_size=config.vocab_size,
+        )
+
+        if config.predict_path:
+            self.path_head = PathPredictionHead(d_model=config.d_model)
+        else:
+            self.path_head = None
+
+        # Length prediction head (predicts word length from path)
+        # Max length is max_char_len (including EOS)
+        self.length_head = LengthPredictionHead(
+            d_model=config.d_model,
+            max_length=config.max_char_len,
+        )
+
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        path_coords: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+        output_hidden_states: bool | None = None,
+    ):
+        """
+        Forward pass of the model.
+
+        Args:
+            path_coords (torch.Tensor): Path coordinates [batch, path_len, 3]
+            input_ids (torch.Tensor): Character token IDs [batch, char_len]
+            attention_mask (torch.Tensor, optional): Attention mask [batch, seq_len]
+            labels (torch.Tensor, optional): Labels for loss calculation [batch, char_len]
+            return_dict (bool, optional): Whether to return ModelOutput object
+            output_hidden_states (bool, optional): Whether to output hidden states
+
+        Returns:
+            SwipeTransformerOutput or tuple: Model outputs with:
+                - loss: Optional loss value
+                - char_logits: Character prediction logits [batch, seq_len, vocab_size]
+                - path_logits: Path prediction logits [batch, seq_len, 3] (if predict_path=True)
+                - length_logits: Length prediction logits [batch, max_length]
+                - last_hidden_state: Hidden states [batch, seq_len, d_model]
+                - pooler_output: SEP token embeddings [batch, d_model] for similarity/embedding tasks
+                - hidden_states: Tuple of hidden states (if output_hidden_states=True)
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = path_coords.shape[0]
+        device = path_coords.device
+
+        # Create [CLS] and [SEP] tokens
+        cls_token = torch.full(
+            (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
+        )
+
+        # 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
+
+        # Encode (batch_first=True is set in TransformerEncoderLayer)
+        hidden_states = self.encoder(embeddings, src_key_padding_mask=src_key_padding_mask)
+
+        # Character prediction
+        char_logits = self.char_head(hidden_states)
+
+        # Path prediction (if enabled)
+        path_logits = None
+        if self.path_head is not None:
+            path_logits = self.path_head(hidden_states)
+
+        # Length prediction from CLS token
+        cls_hidden = hidden_states[:, 0, :]  # [batch, d_model] - CLS at position 0
+        length_logits = self.length_head(cls_hidden)  # [batch, max_length]
+
+        # Extract SEP token embedding for pooler output (embeddings/similarity tasks)
+        # SEP is at position 1 + path_len
+        path_len = path_coords.shape[1]
+        sep_position = 1 + path_len
+        pooler_output = hidden_states[:, sep_position, :]  # [batch, d_model]
+
+        # Compute loss if labels provided
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
+            # Extract character positions from hidden states
+            # Sequence is: [CLS] + path + [SEP] + chars
+            char_start = 1 + path_len + 1  # After [CLS], path, and [SEP]
+            char_hidden = hidden_states[:, char_start : char_start + labels.shape[1], :]
+            char_pred = self.char_head(char_hidden)
+            loss = loss_fct(char_pred.reshape(-1, self.config.vocab_size), labels.reshape(-1))
+
+        if not return_dict:
+            output = (hidden_states, char_logits, length_logits, pooler_output)
+            if path_logits is not None:
+                output = output + (path_logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return SwipeTransformerOutput(
+            loss=loss,
+            char_logits=char_logits,
+            path_logits=path_logits,
+            length_logits=length_logits,
+            last_hidden_state=hidden_states,
+            pooler_output=pooler_output,
+            hidden_states=(hidden_states,) if output_hidden_states else None,
+        )
+
+
+class SwipeCrossEncoderForSequenceClassification(SwipeTransformerPreTrainedModel):
+    """
+    HuggingFace-compatible cross-encoder for sequence classification.
+
+    This model is designed for similarity scoring between swipe paths and words.
+    It extracts the SEP token embedding and passes it through a classification head.
+
+    Args:
+        config (SwipeCrossEncoderConfig): Model configuration
+    """
+
+    config_class = SwipeCrossEncoderConfig
+    base_model_prefix = "swipe_cross_encoder"
+
+    def __init__(self, config: SwipeCrossEncoderConfig):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = config.num_labels
+
+        # Import existing components
+        from .embeddings import MixedEmbedding
+        from .heads import ClassificationHead
+
+        # Embeddings
+        self.embeddings = MixedEmbedding(
+            vocab_size=config.vocab_size,
+            max_path_len=config.max_path_len,
+            max_char_len=config.max_char_len,
+            d_model=config.d_model,
+            dropout=config.dropout,
+        )
+
+        # Transformer encoder
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.n_heads,
+            dim_feedforward=config.d_ff,
+            dropout=config.dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,  # Pre-LayerNorm
+        )
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=config.n_layers,
+            enable_nested_tensor=False,
+        )
+
+        # Classification head
+        self.classifier = ClassificationHead(
+            d_model=config.d_model,
+            num_labels=config.num_labels,
+        )
+
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        path_coords: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+    ):
+        """
+        Forward pass for cross-encoder.
+
+        Args:
+            path_coords (torch.Tensor): Path coordinates [batch, path_len, 3]
+            input_ids (torch.Tensor): Character token IDs [batch, char_len]
+            attention_mask (torch.Tensor, optional): Attention mask [batch, seq_len]
+            labels (torch.Tensor, optional): Labels for loss calculation [batch, num_labels]
+            return_dict (bool, optional): Whether to return ModelOutput object
+
+        Returns:
+            SequenceClassifierOutput or tuple: Model outputs with logits and optional loss
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = path_coords.shape[0]
+        device = path_coords.device
+
+        # Create [CLS] and [SEP] tokens
+        cls_token = torch.full(
+            (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
+        )
+
+        # Get embeddings
+        embeddings = self.embeddings(path_coords, input_ids, cls_token, sep_token)
+
+        # Prepare attention mask
+        if attention_mask is not None:
+            src_key_padding_mask = attention_mask == 0
+        else:
+            src_key_padding_mask = None
+
+        # Encode (batch_first=True is set in TransformerEncoderLayer)
+        hidden_states = self.encoder(embeddings, src_key_padding_mask=src_key_padding_mask)
+
+        # Extract SEP token embedding
+        # SEP is at position 1 + path_len
+        path_len = path_coords.shape[1]
+        sep_position = 1 + path_len
+        sep_embedding = hidden_states[:, sep_position, :]  # [batch, d_model]
+
+        # Classification
+        logits = self.classifier(sep_embedding)  # [batch, num_labels]
+
+        # Compute loss if labels provided
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                else:
+                    self.config.problem_type = "single_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = nn.MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = nn.CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + (hidden_states,)
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=(hidden_states,),
+        )
+
+
+class SwipeModel(SwipeTransformerPreTrainedModel):
+    """
+    Base Swipe model for extracting embeddings.
+
+    .. deprecated::
+        This class is deprecated. Use SwipeTransformerModel instead, which now
+        includes pooler_output for embeddings alongside prediction heads.
+        SwipeTransformerModel provides both predictions AND embeddings in a single model.
+
+    This model returns the SEP token embedding, which can be used for:
+    - Vector databases
+    - Semantic search
+    - Similarity computation
+
+    The SEP token embedding represents the joint encoding of the path and text.
+
+    Usage (Deprecated - use SwipeTransformerModel instead):
+        ```python
+        from transformers import AutoModel
+
+        model = AutoModel.from_pretrained(
+            "your-username/swipe-model",
+            trust_remote_code=True
+        )
+
+        # Get embeddings
+        outputs = model(path_coords=paths, input_ids=tokens)
+        embeddings = outputs.pooler_output  # SEP token embeddings
+        ```
+
+    Args:
+        config (SwipeTransformerConfig or SwipeCrossEncoderConfig): Model configuration
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        # Import existing components
+        from .embeddings import MixedEmbedding
+
+        # Embeddings
+        self.embeddings = MixedEmbedding(
+            vocab_size=config.vocab_size,
+            max_path_len=config.max_path_len,
+            max_char_len=config.max_char_len,
+            d_model=config.d_model,
+            dropout=config.dropout,
+        )
+
+        # Transformer encoder
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.n_heads,
+            dim_feedforward=config.d_ff,
+            dropout=config.dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,  # Pre-LayerNorm
+        )
+        self.encoder = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=config.n_layers,
+            enable_nested_tensor=False,
+        )
+
+        # Initialize weights
+        self.post_init()
+
+    def forward(
+        self,
+        path_coords: torch.Tensor,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+        output_hidden_states: bool | None = None,
+    ):
+        """
+        Forward pass that returns embeddings.
+
+        Args:
+            path_coords (torch.Tensor): Path coordinates [batch, path_len, 3]
+            input_ids (torch.Tensor): Character token IDs [batch, char_len]
+            attention_mask (torch.Tensor, optional): Attention mask [batch, seq_len]
+            return_dict (bool, optional): Whether to return ModelOutput object
+            output_hidden_states (bool, optional): Whether to output all hidden states
+
+        Returns:
+            BaseModelOutputWithPooling with:
+                - last_hidden_state: Full sequence hidden states [batch, seq_len, d_model]
+                - pooler_output: SEP token embeddings [batch, d_model]
+                - hidden_states: Tuple of hidden states (if output_hidden_states=True)
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        batch_size = path_coords.shape[0]
+        device = path_coords.device
+
+        # Create [CLS] and [SEP] tokens
+        cls_token = torch.full(
+            (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
+        )
+
+        # Get embeddings
+        embeddings = self.embeddings(path_coords, input_ids, cls_token, sep_token)
+
+        # Prepare attention mask
+        if attention_mask is not None:
+            src_key_padding_mask = attention_mask == 0
+        else:
+            src_key_padding_mask = None
+
+        # Encode (batch_first=True is set in TransformerEncoderLayer)
+        hidden_states = self.encoder(embeddings, src_key_padding_mask=src_key_padding_mask)
+
+        # Extract SEP token embedding (pooler output)
+        # SEP is at position 1 + path_len
+        path_len = path_coords.shape[1]
+        sep_position = 1 + path_len
+        pooler_output = hidden_states[:, sep_position, :]  # [batch, d_model]
+
+        if not return_dict:
+            return (hidden_states, pooler_output)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=hidden_states,
+            pooler_output=pooler_output,
+            hidden_states=(hidden_states,) if output_hidden_states else None,
+        )

processing_swipe.py

@@ -0,0 +1,263 @@
+"""Processor for handling multimodal swipe inputs (path + text)."""
+
+import numpy as np
+import torch
+from transformers import ProcessorMixin
+
+
+class SwipeProcessor(ProcessorMixin):
+    """
+    Processor for handling multimodal swipe inputs (path coordinates + text).
+
+    This processor combines path coordinate preprocessing with text tokenization,
+    creating the inputs needed for SwipeTransformer models.
+
+    Args:
+        tokenizer: SwipeTokenizer instance
+        max_path_len (int): Maximum path length. Defaults to 64.
+        max_char_len (int): Maximum character length. Defaults to 38.
+    """
+
+    attributes = ["tokenizer"]
+    tokenizer_class = "AutoTokenizer"  # Will use auto_map from tokenizer_config.json
+
+    def __init__(self, tokenizer=None, max_path_len: int = 64, max_char_len: int = 38):
+        self.tokenizer = tokenizer
+        self.max_path_len = max_path_len
+        self.max_char_len = max_char_len
+        # 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,
+        path_coords: list[list[list[float]]] | torch.Tensor | np.ndarray | None = None,
+        text: str | list[str] | None = None,
+        padding: bool | str = True,
+        truncation: bool = True,
+        max_length: int | None = None,
+        return_tensors: str | None = "pt",
+        **kwargs,
+    ):
+        """
+        Process path coordinates and text into model inputs.
+
+        Args:
+            path_coords: List of paths or tensor [batch, path_len, 3]
+                        Each point is (x, y, time). Can be None if only processing text.
+            text: String or list of strings to encode. Can be None if only processing paths.
+            padding: Whether to pad sequences. Can be True/False or "max_length"
+            truncation: Whether to truncate sequences
+            max_length: Maximum sequence length for text (overrides max_char_len)
+            return_tensors: "pt" for PyTorch, "np" for NumPy, None for lists
+            **kwargs: Additional keyword arguments
+
+        Returns:
+            Dictionary with:
+                - path_coords: [batch, max_path_len, 3] (if path_coords provided)
+                - input_ids: [batch, max_char_len] (if text provided)
+                - attention_mask: [batch, total_seq_len]
+        """
+        if path_coords is None and text is None:
+            raise ValueError("Must provide either path_coords or text (or both)")
+
+        # Determine batch size
+        if path_coords is not None:
+            # Handle path coordinates
+            if isinstance(path_coords, (list, tuple)):
+                # Check if it's a batch or single path
+                if len(path_coords) > 0 and isinstance(path_coords[0][0], (list, tuple)):
+                    # Batch of paths [[path1], [path2], ...]
+                    path_coords = torch.tensor(path_coords, dtype=torch.float32)
+                else:
+                    # Single path [[x,y,t], [x,y,t], ...]
+                    path_coords = torch.tensor([path_coords], dtype=torch.float32)
+            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)
+            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]
+        elif text is not None:
+            if isinstance(text, str):
+                batch_size = 1
+                text = [text]
+            else:
+                batch_size = len(text)
+        else:
+            batch_size = 1
+
+        result = {}
+
+        # Process path coordinates
+        if path_coords is not None:
+            current_path_len = path_coords.shape[1]
+
+            # Truncate if needed
+            if truncation and current_path_len > self.max_path_len:
+                path_coords = path_coords[:, : self.max_path_len, :]
+                current_path_len = self.max_path_len
+
+            # Pad if needed
+            if padding and current_path_len < self.max_path_len:
+                pad_len = self.max_path_len - current_path_len
+                path_coords = torch.cat([path_coords, torch.zeros(batch_size, pad_len, 3)], dim=1)
+
+            # Create path mask (1 = real data, 0 = padding)
+            path_mask = torch.ones(batch_size, self.max_path_len, dtype=torch.long)
+            if padding and current_path_len < self.max_path_len:
+                path_mask[:, current_path_len:] = 0
+
+            result["path_coords"] = path_coords
+            # Store path_mask internally for attention_mask construction
+            _path_mask = path_mask
+        else:
+            # No path coords provided, create empty/zero tensors
+            path_coords = torch.zeros(batch_size, self.max_path_len, 3)
+            _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,
+            )
+
+            # 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
+            if return_tensors == "pt":
+                char_tokens = torch.full(
+                    (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)
+            elif return_tensors == "np":
+                char_tokens = np.full(
+                    (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)
+            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)]
+
+            result["input_ids"] = char_tokens
+
+        # Create combined attention mask: [CLS] + path + [SEP] + chars
+        # Sequence structure: [CLS:1] + _path_mask + [SEP:1] + _char_mask
+        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":
+            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
+        if return_tensors == "np":
+            for key in result:
+                if isinstance(result[key], torch.Tensor):
+                    result[key] = result[key].numpy()
+        elif return_tensors is None:
+            for key in result:
+                if isinstance(result[key], torch.Tensor):
+                    result[key] = result[key].tolist()
+
+        return result
+
+    def batch_decode(self, token_ids, **kwargs):
+        """
+        Decode token IDs to strings.
+
+        Args:
+            token_ids: Token IDs to decode
+            **kwargs: Additional arguments passed to tokenizer
+
+        Returns:
+            List of decoded strings
+        """
+        return self.tokenizer.batch_decode(token_ids, **kwargs)
+
+    def decode(self, token_ids, **kwargs):
+        """
+        Decode single sequence of token IDs to string.
+
+        Args:
+            token_ids: Token IDs to decode
+            **kwargs: Additional arguments passed to tokenizer
+
+        Returns:
+            Decoded string
+        """
+        return self.tokenizer.decode(token_ids, **kwargs)

processor_config.json

@@ -0,0 +1,8 @@
+{
+  "max_char_len": 48,
+  "max_path_len": 128,
+  "processor_class": "SwipeProcessor",
+  "auto_map": {
+    "AutoProcessor": "processing_swipe.SwipeProcessor"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,8 @@
+{
+  "cls_token": "[CLS]",
+  "eos_token": "[EOS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenization_swipe.py

@@ -0,0 +1,243 @@
+"""HuggingFace-compatible tokenizer for SwipeTransformer."""
+
+import json
+import os
+
+from transformers import PreTrainedTokenizer
+
+from .tokenizer import CharacterTokenizer
+
+
+class SwipeTokenizer(PreTrainedTokenizer):
+    """
+    HuggingFace-compatible tokenizer that wraps the existing CharacterTokenizer.
+
+    This tokenizer provides a HuggingFace-compatible interface for the custom
+    character-level tokenization used in the swipe keyboard model.
+
+    Args:
+        vocab_file (str, optional): Path to vocabulary file
+        unk_token (str): Unknown token. Defaults to "[UNK]"
+        sep_token (str): Separator token. Defaults to "[SEP]"
+        pad_token (str): Padding token. Defaults to "[PAD]"
+        cls_token (str): Classification token. Defaults to "[CLS]"
+        mask_token (str): Mask token. Defaults to "[MASK]"
+        eos_token (str): End-of-sequence token. Defaults to "[EOS]"
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str | None = None,
+        unk_token: str = "[UNK]",
+        sep_token: str = "[SEP]",
+        pad_token: str = "[PAD]",
+        cls_token: str = "[CLS]",
+        mask_token: str = "[MASK]",
+        eos_token: str = "[EOS]",
+        **kwargs,
+    ):
+        # Initialize internal CharacterTokenizer BEFORE calling super().__init__()
+        # because super().__init__() will call get_vocab() which needs self._tokenizer
+        if vocab_file is not None and os.path.exists(vocab_file):
+            # Load from vocab file
+            with open(vocab_file, encoding="utf-8") as f:
+                vocab_data = json.load(f)
+
+            # Extract vocabulary (excluding ALL special tokens)
+            # All special tokens that should NOT be passed to CharacterTokenizer
+            # Convert AddedToken objects to strings
+            special_tokens_to_exclude = {
+                str(pad_token),
+                str(cls_token),
+                str(sep_token),
+                str(mask_token),
+                str(unk_token),
+                str(eos_token),
+                "[PUNC]",
+            }
+
+            if "chars" in vocab_data:
+                # Filter out special tokens from the chars list
+                vocab = set(c for c in vocab_data["chars"] if c not in special_tokens_to_exclude)
+            elif "char_to_id" in vocab_data:
+                # Get all characters except special tokens
+                vocab = set(
+                    c for c in vocab_data["char_to_id"].keys() if c not in special_tokens_to_exclude
+                )
+            else:
+                vocab = None
+
+            self._tokenizer = CharacterTokenizer(vocab=vocab)
+        else:
+            # Default vocab (will be built from dataset during conversion)
+            self._tokenizer = CharacterTokenizer()
+
+        super().__init__(
+            unk_token=unk_token,
+            sep_token=sep_token,
+            pad_token=pad_token,
+            cls_token=cls_token,
+            mask_token=mask_token,
+            eos_token=eos_token,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        """Return the size of the vocabulary"""
+        return self._tokenizer.vocab_size
+
+    def get_vocab(self):
+        """Return the vocabulary as a dict"""
+        return self._tokenizer.char_to_id.copy()
+
+    def _tokenize(self, text: str) -> list[str]:
+        """
+        Tokenize a string into tokens (characters).
+
+        Args:
+            text (str): Text to tokenize
+
+        Returns:
+            List[str]: List of character tokens
+        """
+        # Convert to lowercase and split into characters
+        return list(text.lower())
+
+    def _convert_token_to_id(self, token: str) -> int:
+        """
+        Convert a token (character) to an id using the vocabulary.
+
+        Args:
+            token (str): Token to convert
+
+        Returns:
+            int: Token ID
+        """
+        return self._tokenizer.char_to_id.get(token, self._tokenizer.unk_token_id)
+
+    def _convert_id_to_token(self, index: int) -> str:
+        """
+        Convert an index to a token using the vocabulary.
+
+        Args:
+            index (int): Token ID
+
+        Returns:
+            str: Token (character)
+        """
+        return self._tokenizer.id_to_char.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        """
+        Convert a list of tokens (characters) to a string.
+
+        Args:
+            tokens (List[str]): List of tokens
+
+        Returns:
+            str: Concatenated string
+        """
+        # Filter out special tokens
+        special_tokens = {
+            self.pad_token,
+            self.cls_token,
+            self.sep_token,
+            self.mask_token,
+            self.unk_token,
+            self.eos_token,
+        }
+        filtered = [t for t in tokens if t not in special_tokens]
+        return "".join(filtered)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple:
+        """
+        Save the tokenizer vocabulary to a directory.
+
+        Args:
+            save_directory (str): Directory to save the vocabulary
+            filename_prefix (str, optional): Optional prefix for the vocabulary file
+
+        Returns:
+            tuple: Tuple containing the path to the saved vocabulary file
+        """
+        if not os.path.isdir(save_directory):
+            os.makedirs(save_directory, exist_ok=True)
+
+        vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        )
+
+        # Save vocabulary and mappings
+        vocab_data = {
+            "chars": sorted(list(set(self._tokenizer.char_to_id.keys()))),
+            "char_to_id": self._tokenizer.char_to_id,
+            "special_tokens": {
+                "pad_token": self.pad_token,
+                "cls_token": self.cls_token,
+                "sep_token": self.sep_token,
+                "mask_token": self.mask_token,
+                "unk_token": self.unk_token,
+                "eos_token": self.eos_token,
+            },
+        }
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            json.dump(vocab_data, f, ensure_ascii=False, indent=2)
+
+        return (vocab_file,)
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+    ) -> list[int]:
+        """
+        Build model inputs from a sequence by adding special tokens.
+
+        For swipe models, we don't add special tokens here as they are
+        handled separately (CLS and SEP are managed by the model/processor).
+
+        Args:
+            token_ids_0 (List[int]): First sequence
+            token_ids_1 (List[int], optional): Second sequence
+
+        Returns:
+            List[int]: Sequence with special tokens
+        """
+        # For swipe models, special tokens are handled by the processor
+        # Just return the tokens as-is
+        if token_ids_1 is None:
+            return token_ids_0
+        return token_ids_0 + token_ids_1
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        """
+        Retrieve sequence ids from a token list.
+
+        Args:
+            token_ids_0 (List[int]): First sequence
+            token_ids_1 (List[int], optional): Second sequence
+            already_has_special_tokens (bool): Whether tokens already have special tokens
+
+        Returns:
+            List[int]: Mask (1 for special tokens, 0 for normal tokens)
+        """
+        # All special token handling is done by the processor
+        # Return all zeros
+        if already_has_special_tokens:
+            if token_ids_1 is not None:
+                raise ValueError(
+                    "You should not supply a second sequence if the provided sequence already has special tokens."
+                )
+            return [0] * len(token_ids_0)
+
+        if token_ids_1 is None:
+            return [0] * len(token_ids_0)
+        return [0] * len(token_ids_0) + [0] * len(token_ids_1)

tokenizer.py

@@ -0,0 +1,170 @@
+"""Character tokenizer and vocabulary utilities for swipe keyboard dataset."""
+
+import hashlib
+
+import torch
+
+
+class CharacterTokenizer:
+    """Character-level tokenizer for swipe keyboard words."""
+
+    def __init__(self, vocab: set | None = None):
+        """
+        Initialize tokenizer with vocabulary.
+
+        Args:
+            vocab: Optional set of characters. If None, will use printable ASCII.
+        """
+        # Special tokens
+        self.pad_token = "[PAD]"
+        self.cls_token = "[CLS]"
+        self.sep_token = "[SEP]"
+        self.mask_token = "[MASK]"
+        self.unk_token = "[UNK]"
+        self.eos_token = "[EOS]"  # End of word token
+        self.punc_token = "[PUNC]"
+
+        self.special_tokens = [
+            self.pad_token,  # 0
+            self.cls_token,  # 1
+            self.sep_token,  # 2
+            self.mask_token,  # 3
+            self.unk_token,  # 4
+            self.eos_token,  # 5
+            self.punc_token,  # 6
+        ]
+
+        # Build vocabulary deterministically (lowercase letters + digits).
+        chars = set(chr(i) for i in range(ord("a"), ord("z") + 1))
+        chars.update(str(d) for d in range(10))
+        if vocab is not None:
+            # Allow explicit extension for special cases
+            chars.update(vocab)
+
+        self.char_to_id = {token: idx for idx, token in enumerate(self.special_tokens)}
+        for idx, char in enumerate(sorted(chars), start=len(self.special_tokens)):
+            self.char_to_id[char] = idx
+
+        self.id_to_char = {idx: char for char, idx in self.char_to_id.items()}
+        self.vocab_size = len(self.char_to_id)
+
+    def encode(self, text: str) -> list[int]:
+        """Encode text to token IDs (case-insensitive, punctuation -> [PUNC])."""
+        unk_id = self.char_to_id[self.unk_token]
+        punc_id = self.char_to_id[self.punc_token]
+        tokens = []
+        for char in text.lower():
+            if char.isalpha() or char.isdigit():
+                tokens.append(self.char_to_id.get(char, unk_id))
+            else:
+                tokens.append(punc_id)
+        return tokens
+
+    def decode(self, token_ids: list[int]) -> str:
+        """Decode token IDs to text, stopping at EOS token."""
+        chars = []
+        for token_id in token_ids:
+            if token_id in self.id_to_char:
+                char = self.id_to_char[token_id]
+                # Stop at EOS token
+                if char == self.eos_token:
+                    break
+                # Skip other special tokens except for debugging
+                if char not in self.special_tokens or char == " ":
+                    chars.append(char)
+        return "".join(chars)
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.char_to_id[self.pad_token]
+
+    @property
+    def cls_token_id(self) -> int:
+        return self.char_to_id[self.cls_token]
+
+    @property
+    def sep_token_id(self) -> int:
+        return self.char_to_id[self.sep_token]
+
+    @property
+    def mask_token_id(self) -> int:
+        return self.char_to_id[self.mask_token]
+
+    @property
+    def unk_token_id(self) -> int:
+        return self.char_to_id[self.unk_token]
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.char_to_id[self.eos_token]
+
+    @property
+    def punc_token_id(self) -> int:
+        return self.char_to_id[self.punc_token]
+
+
+def vocab_hash(tokenizer: CharacterTokenizer) -> str:
+    """Stable hash of the tokenizer's id->token mapping (includes specials)."""
+    ordered_tokens = [tokenizer.id_to_char[i] for i in range(tokenizer.vocab_size)]
+    joined = "\n".join(ordered_tokens).encode("utf-8")
+    return hashlib.sha256(joined).hexdigest()
+
+
+def compute_char_frequency_weights(
+    tokenizer: CharacterTokenizer,
+    dataset,
+    max_samples: int | None = None,
+    weight_exponent: float = 1.0,
+):
+    """Compute inverse log frequency weights for characters.
+
+    Args:
+        tokenizer: CharacterTokenizer used for encoding
+        dataset: HF dataset or iterable of samples with a 'word' field
+        max_samples: Optional cap on samples to scan
+
+    Returns:
+        torch.Tensor of shape [vocab_size] with weights normalized to mean=1.
+        Padding token weight is set to the non-pad mean (not zero) so min>0.
+    """
+    counts = torch.ones(tokenizer.vocab_size, dtype=torch.float)  # start at 1 for smoothing
+
+    # Collect all token IDs first for vectorized counting
+    all_token_ids = []
+    for idx, sample in enumerate(dataset):
+        if max_samples is not None and idx >= max_samples:
+            break
+
+        # Encode lowercase characters and append EOS (matches training labels)
+        token_ids = tokenizer.encode(sample["word"]) + [tokenizer.eos_token_id]
+        all_token_ids.extend(token_ids)
+
+    # Use bincount for efficient vectorized counting
+    if all_token_ids:
+        token_tensor = torch.tensor(all_token_ids, dtype=torch.long)
+        bincount_result = torch.bincount(token_tensor, minlength=tokenizer.vocab_size).float()
+        counts = counts + bincount_result
+
+    # Padding is never a supervised label, but keep a finite weight
+    pad_id = tokenizer.pad_token_id
+    counts[pad_id] = counts[pad_id]  # leave smoothing value as-is
+
+    # Inverse log weighting; add 1 inside log to avoid div-by-zero
+    weights = 1.0 / torch.log1p(counts)
+
+    # Use non-pad mean for pad token to avoid zero/inf
+    non_pad_mask = torch.ones_like(weights, dtype=torch.bool)
+    non_pad_mask[pad_id] = False
+    non_pad_mean = weights[non_pad_mask].mean().clamp_min(1e-8)
+    weights[pad_id] = non_pad_mean
+
+    # Optional tempering (e.g., exponent <1 flattens extremes)
+    if weight_exponent != 1.0:
+        weights = torch.pow(weights, weight_exponent)
+
+    # Normalize to keep loss scale stable (mean of non-pad tokens -> 1)
+    non_pad_mean = weights[non_pad_mask].mean().clamp_min(1e-8)
+    weights[pad_id] = non_pad_mean
+    weights = weights / non_pad_mean
+
+    return weights

tokenizer_config.json

@@ -0,0 +1,69 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "[EOS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "eos_token": "[EOS]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "processor_class": "SwipeProcessor",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "SwipeTokenizer",
+  "unk_token": "[UNK]",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_swipe.SwipeTokenizer",
+      null
+    ]
+  }
+}

vocab.json

@@ -0,0 +1,100 @@
+{
+  "chars": [
+    "0",
+    "1",
+    "2",
+    "3",
+    "4",
+    "5",
+    "6",
+    "7",
+    "8",
+    "9",
+    "[CLS]",
+    "[EOS]",
+    "[MASK]",
+    "[PAD]",
+    "[PUNC]",
+    "[SEP]",
+    "[UNK]",
+    "a",
+    "b",
+    "c",
+    "d",
+    "e",
+    "f",
+    "g",
+    "h",
+    "i",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "o",
+    "p",
+    "q",
+    "r",
+    "s",
+    "t",
+    "u",
+    "v",
+    "w",
+    "x",
+    "y",
+    "z"
+  ],
+  "char_to_id": {
+    "[PAD]": 0,
+    "[CLS]": 1,
+    "[SEP]": 2,
+    "[MASK]": 3,
+    "[UNK]": 4,
+    "[EOS]": 5,
+    "[PUNC]": 6,
+    "0": 7,
+    "1": 8,
+    "2": 9,
+    "3": 10,
+    "4": 11,
+    "5": 12,
+    "6": 13,
+    "7": 14,
+    "8": 15,
+    "9": 16,
+    "a": 17,
+    "b": 18,
+    "c": 19,
+    "d": 20,
+    "e": 21,
+    "f": 22,
+    "g": 23,
+    "h": 24,
+    "i": 25,
+    "j": 26,
+    "k": 27,
+    "l": 28,
+    "m": 29,
+    "n": 30,
+    "o": 31,
+    "p": 32,
+    "q": 33,
+    "r": 34,
+    "s": 35,
+    "t": 36,
+    "u": 37,
+    "v": 38,
+    "w": 39,
+    "x": 40,
+    "y": 41,
+    "z": 42
+  },
+  "special_tokens": {
+    "pad_token": "[PAD]",
+    "cls_token": "[CLS]",
+    "sep_token": "[SEP]",
+    "mask_token": "[MASK]",
+    "unk_token": "[UNK]",
+    "eos_token": "[EOS]"
+  }
+}