Add operator-aware splitting for word number extraction

- SYMBOL_OP_TOKENS for ' +', ' -', etc.
- WORD_OP_TOKENS for 'plus', 'minus', etc. (token IDs)
- ALL_OP_TOKENS combines both for unified lookup
- Separate a_pool and b_pool attention modules
- _find_op_position() to locate operator in sequence
- Split hidden states at operator: a tokens before, b tokens after
- Each operand pooled separately before digit prediction

.gitignore

@@ -1,2 +1,3 @@
 __pycache__/
 *.pyc
+.pt file

llm_integration/model.py

@@ -757,7 +757,7 @@ class HybridExtractor(nn.Module):
     """
 
     DIGIT_TOKENS = set(range(32, 42))
-    OPERATOR_TOKENS = {
+    SYMBOL_OP_TOKENS = {
         1232: 0,   # ' +' -> add
         731: 1,    # ' -' -> sub
         1672: 2,   # ' *' -> mul
@@ -766,15 +766,23 @@ class HybridExtractor(nn.Module):
         1758: 5,   # ' ==' -> eq
     }
     WORD_OP_TOKENS = {
-        'plus': 0, 'minus': 1, 'times': 2,
-        'greater': 3, 'less': 4, 'equals': 5, 'equal': 5,
+        2068: 0,   # 'plus' -> add
+        8500: 1,   # 'minus' -> sub
+        1580: 2,   # 'times' -> mul
+        6301: 3,   # 'greater' -> gt
+        1912: 4,   # 'less' -> lt
+        16364: 5,  # 'equals' -> eq
+        11540: 5,  # 'equal' -> eq
     }
+    ALL_OP_TOKENS = {**SYMBOL_OP_TOKENS, **WORD_OP_TOKENS}
 
     def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
         super().__init__()
         self.hidden_dim = hidden_dim
 
         self.attention_pool = AttentionPooling(hidden_dim, num_heads)
+        self.a_pool = AttentionPooling(hidden_dim, num_heads)
+        self.b_pool = AttentionPooling(hidden_dim, num_heads)
 
         self.a_digit_pred = nn.Sequential(
             nn.Linear(hidden_dim, intermediate_dim),
@@ -803,9 +811,18 @@ class HybridExtractor(nn.Module):
                 return True
         return False
 
+    def _find_op_position(self, token_ids: torch.Tensor) -> int:
+        """Find position of operator token, returns -1 if not found."""
+        tokens = token_ids.tolist()
+        for i, tid in enumerate(tokens):
+            if tid in self.ALL_OP_TOKENS:
+                return i
+        return -1
+
     def _extract_from_digits(self, token_ids: torch.Tensor) -> tuple:
         """
         Extract values directly from digit tokens (hardcoded lookup).
+        Handles both symbol operators (' +') and word operators ('plus').
         Returns (a_value, b_value, op_idx) or None if pattern not found.
         """
         tokens = token_ids.tolist()
@@ -813,9 +830,9 @@ class HybridExtractor(nn.Module):
         op_pos = -1
         op_idx = 0
         for i, tid in enumerate(tokens):
-            if tid in self.OPERATOR_TOKENS:
+            if tid in self.ALL_OP_TOKENS:
                 op_pos = i
-                op_idx = self.OPERATOR_TOKENS[tid]
+                op_idx = self.ALL_OP_TOKENS[tid]
                 break
 
         if op_pos == -1:
@@ -934,11 +951,34 @@ class HybridExtractor(nn.Module):
                 a_digit_logits_list.append(None)
                 b_digit_logits_list.append(None)
             else:
-                sample_pooled = pooled[i]
+                sample_hidden = hidden[i:i+1]
+                sample_mask = mask[i:i+1]
+
+                seq_mask = mask[i].bool()
+                valid_len = int(seq_mask.sum().item())
+                start_pos = hidden.shape[1] - valid_len
+                valid_tokens = token_ids[i, start_pos:] if token_ids is not None else None
+
+                op_pos = self._find_op_position(valid_tokens) if valid_tokens is not None else -1
+
+                if op_pos > 0 and op_pos < valid_len - 1:
+                    a_end = start_pos + op_pos
+                    b_start = start_pos + op_pos + 1
+
+                    a_mask = torch.zeros_like(sample_mask)
+                    a_mask[0, start_pos:a_end] = 1.0
+                    b_mask = torch.zeros_like(sample_mask)
+                    b_mask[0, b_start:] = sample_mask[0, b_start:]
+
+                    a_pooled = self.a_pool(sample_hidden, a_mask)[0]
+                    b_pooled = self.b_pool(sample_hidden, b_mask)[0]
+                else:
+                    a_pooled = pooled[i]
+                    b_pooled = pooled[i]
 
-                a_digit_logits = self.a_digit_pred(sample_pooled)
-                b_digit_logits = self.b_digit_pred(sample_pooled)
-                op_logits = self.op_predictor(sample_pooled)
+                a_digit_logits = self.a_digit_pred(a_pooled)
+                b_digit_logits = self.b_digit_pred(b_pooled)
+                op_logits = self.op_predictor(pooled[i])
 
                 a_val, a_bits = self._digits_to_value_and_bits(a_digit_logits, device)
                 b_val, b_bits = self._digits_to_value_and_bits(b_digit_logits, device)