Next-Action Prediction

import orbax.checkpoint as ocp

def compute_loss(model, states, actions, targets):
  """Cross-entropy loss for action prediction."""
  logits = model(states, actions)

  # Flatten for cross-entropy
  logits_flat = logits.reshape(-1, logits.shape[-1])
  targets_flat = targets.reshape(-1)

  loss = optax.softmax_cross_entropy_with_integer_labels(
    logits_flat, targets_flat
  ).mean()

  return loss


@nnx.jit
def train_step(model, optimizer, states, actions, targets):
  """Single training step with gradient update."""

  def loss_fn(model):
    return compute_loss(model, states, actions, targets)

  loss, grads = nnx.value_and_grad(loss_fn)(model)
  optimizer.update(grads)

  return loss


def train(
  model: DecisionTransformer,
  train_dataset,
  config: dict,
):
  """Main training loop."""

  # Optimizer with warmup + cosine decay
  schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=config['lr'],
    warmup_steps=config['warmup_steps'],
    decay_steps=config['total_steps'],
  )
  optimizer = nnx.Optimizer(model, optax.adamw(schedule))

  # Checkpointing
  ckpt_mgr = ocp.CheckpointManager(
    config['checkpoint_dir'],
    options=ocp.CheckpointManagerOptions(max_to_keep=3)
  )

  step = 0
  for epoch in range(config['epochs']):
    for batch in train_dataset:
      states = batch['states']       # [B, T, state_dim]
      actions = batch['actions']     # [B, T]
      targets = batch['next_actions']  # [B, T] (shifted by 1)

      loss = train_step(model, optimizer, states, actions, targets)

      if step % 100 == 0:
        print(f"Step {step}, Loss: {loss:.4f}")
        wandb.log({"loss": loss, "step": step})

      if step % 1000 == 0:
        ckpt_mgr.save(step, args=ocp.args.StandardSave(model))

      step += 1

  return model


# === Run Training ===
config = {
  'state_dim': 512,
  'action_vocab_size': 512,
  'd_model': 512,
  'n_heads': 8,
  'n_layers': 6,
  'max_seq_len': 256,
  'lr': 1e-4,
  'warmup_steps': 1000,
  'total_steps': 100000,
  'epochs': 10,
  'batch_size': 32,
  'checkpoint_dir': './checkpoints',
}

model = create_model(config)
trained_model = train(model, train_dataset, config)

import json
import numpy as np
from pathlib import Path
from sentence_transformers import SentenceTransformer
from dataclasses import dataclass
from typing import List, Tuple

@dataclass
class ProcessedSequence:
  states: np.ndarray      # [T, state_dim]
  actions: np.ndarray     # [T]
  next_actions: np.ndarray  # [T]

class DatasetBuilder:
  def __init__(self, raw_data_dir: str, output_dir: str):
    self.raw_data_dir = Path(raw_data_dir)
    self.output_dir = Path(output_dir)
    self.output_dir.mkdir(exist_ok=True)

    # Encoders
    self.text_encoder = SentenceTransformer('all-MiniLM-L6-v2')
    self.app_vocab = {}  # app_id -> int
    self.action_vocab = self._build_action_vocab()

  def _build_action_vocab(self) -> dict:
    """Build action vocabulary."""
    actions = [
      'SWITCH_APP', 'LAUNCH_APP', 'CLOSE_APP',
      'CLICK', 'RIGHT_CLICK', 'DOUBLE_CLICK', 'SCROLL',
      'HOTKEY', 'TYPE_TEXT',
      'NEW_TAB', 'CLOSE_TAB', 'SWITCH_TAB',
      'OPEN_FILE', 'SAVE_FILE',
      'SPOTLIGHT', 'IDLE',
    ]
    return {a: i for i, a in enumerate(actions)}

  def encode_state(self, event: dict) -> np.ndarray:
    """Encode single event to state vector."""
    # App embedding (one-hot or learned)
    app_id = event['appBundleId']
    if app_id not in self.app_vocab:
      self.app_vocab[app_id] = len(self.app_vocab)
    app_idx = self.app_vocab[app_id]
    app_onehot = np.zeros(100)  # Max 100 apps
    app_onehot[min(app_idx, 99)] = 1.0

    # Window title embedding
    title = event.get('windowTitle', '')
    title_emb = self.text_encoder.encode(title)  # [384]

    # Temporal features
    timestamp = event['timestamp']
    hour = (timestamp % 86400) / 86400  # Normalized hour

    # Cursor position (normalized)
    cursor = event.get('cursorPosition', [0, 0])
    cursor_norm = [c / 2000 for c in cursor]  # Assume 2000px max

    # Combine all features
    state = np.concatenate([
      app_onehot,        # [100]
      title_emb,         # [384]
      [hour],            # [1]
      cursor_norm,       # [2]
    ])  # Total: 487 -> pad to 512

    return np.pad(state, (0, 512 - len(state)))

  def encode_action(self, event: dict) -> int:
    """Map event to action vocabulary index."""
    action_type = event['actionType']
    return self.action_vocab.get(action_type, self.action_vocab['IDLE'])

  def build_sequences(
    self,
    events: List[dict],
    seq_len: int = 64
  ) -> List[ProcessedSequence]:
    """Convert raw events to training sequences."""
    sequences = []

    for i in range(0, len(events) - seq_len - 1, seq_len // 2):
      window = events[i:i + seq_len + 1]

      states = np.array([self.encode_state(e) for e in window[:-1]])
      actions = np.array([self.encode_action(e) for e in window[:-1]])
      next_actions = np.array([self.encode_action(e) for e in window[1:]])

      sequences.append(ProcessedSequence(
        states=states,
        actions=actions,
        next_actions=next_actions
      ))

    return sequences

  def process_all(self):
    """Process all raw data files."""
    all_events = []

    for f in self.raw_data_dir.glob('*.json'):
      with open(f) as fp:
        all_events.extend(json.load(fp))

    # Sort by timestamp
    all_events.sort(key=lambda x: x['timestamp'])

    # Build sequences
    sequences = self.build_sequences(all_events)

    # Save as numpy arrays
    np.savez(
      self.output_dir / 'train.npz',
      states=np.array([s.states for s in sequences]),
      actions=np.array([s.actions for s in sequences]),
      next_actions=np.array([s.next_actions for s in sequences]),
    )

    print(f"Saved {len(sequences)} sequences")


# Usage
builder = DatasetBuilder('./raw_events', './processed')
builder.process_all()

#!/usr/bin/env python3
"""
next_action_quickstart.py

Minimal end-to-end pipeline:
1. Load preprocessed data
2. Train Decision Transformer
3. Evaluate predictions
"""

import jax
import jax.numpy as jnp
from flax import nnx
import optax
import numpy as np
from pathlib import Path

# === Config ===
CONFIG = {
  'state_dim': 512,
  'action_vocab_size': 64,  # Start small
  'd_model': 256,
  'n_heads': 4,
  'n_layers': 4,
  'max_seq_len': 64,
  'lr': 3e-4,
  'batch_size': 16,
  'epochs': 5,
  'data_path': './processed/train.npz',
}

# === Load Data ===
def load_data(path: str):
  data = np.load(path)
  return {
    'states': jnp.array(data['states']),
    'actions': jnp.array(data['actions']),
    'next_actions': jnp.array(data['next_actions']),
  }

# === Simple Dataloader ===
def batch_iterator(data, batch_size):
  n = len(data['states'])
  indices = np.random.permutation(n)

  for i in range(0, n - batch_size, batch_size):
    idx = indices[i:i + batch_size]
    yield {
      'states': data['states'][idx],
      'actions': data['actions'][idx],
      'next_actions': data['next_actions'][idx],
    }

# === Training ===
def main():
  print("Loading data...")
  data = load_data(CONFIG['data_path'])
  print(f"Loaded {len(data['states'])} sequences")

  print("Creating model...")
  model = create_model(CONFIG)

  # Count parameters
  param_count = sum(
    x.size for x in jax.tree.leaves(nnx.state(model))
  )
  print(f"Model parameters: {param_count:,}")

  # Optimizer
  optimizer = nnx.Optimizer(model, optax.adam(CONFIG['lr']))

  print("Training...")
  for epoch in range(CONFIG['epochs']):
    total_loss = 0
    n_batches = 0

    for batch in batch_iterator(data, CONFIG['batch_size']):
      loss = train_step(
        model, optimizer,
        batch['states'],
        batch['actions'],
        batch['next_actions']
      )
      total_loss += loss
      n_batches += 1

    avg_loss = total_loss / n_batches
    print(f"Epoch {epoch + 1}/{CONFIG['epochs']}, Loss: {avg_loss:.4f}")

  # === Evaluation ===
  print("\nEvaluating...")
  evaluate(model, data)


def evaluate(model, data, n_samples=100):
  """Compute top-k accuracy."""
  correct_top1 = 0
  correct_top3 = 0

  for i in range(min(n_samples, len(data['states']))):
    states = data['states'][i:i+1]
    actions = data['actions'][i:i+1]
    targets = data['next_actions'][i]

    logits = model(states, actions)
    preds = jnp.argsort(logits[0, -1])[::-1]  # Last position

    target = targets[-1]
    if preds[0] == target:
      correct_top1 += 1
    if target in preds[:3]:
      correct_top3 += 1

  print(f"Top-1 Accuracy: {correct_top1 / n_samples:.2%}")
  print(f"Top-3 Accuracy: {correct_top3 / n_samples:.2%}")


if __name__ == '__main__':
  main()

import zmq
import json
import jax.numpy as jnp
from collections import deque

class ActionPredictor:
  """Real-time next-action prediction service."""

  def __init__(self, model, max_context: int = 64):
    self.model = model
    self.max_context = max_context
    self.state_buffer = deque(maxlen=max_context)
    self.action_buffer = deque(maxlen=max_context)

    # ZMQ for receiving events from Swift logger
    self.context = zmq.Context()
    self.socket = self.context.socket(zmq.SUB)
    self.socket.connect("tcp://localhost:5555")
    self.socket.setsockopt_string(zmq.SUBSCRIBE, "")

  def encode_event(self, event: dict) -> tuple:
    """Convert raw event to (state, action) pair."""
    # Reuse encoding logic from dataset builder
    state = encode_state(event)
    action = encode_action(event)
    return state, action

  def predict_next(self) -> list:
    """Get top-k predictions for next action."""
    if len(self.state_buffer) < 2:
      return []

    # Prepare input tensors
    states = jnp.array([list(self.state_buffer)])
    actions = jnp.array([list(self.action_buffer)])

    # Forward pass
    logits = self.model(states, actions)
    probs = jax.nn.softmax(logits[0, -1])

    # Top 5 predictions
    top_k = 5
    top_indices = jnp.argsort(probs)[::-1][:top_k]

    predictions = [
      {
        'action': ACTION_NAMES[int(idx)],
        'probability': float(probs[idx])
      }
      for idx in top_indices
    ]

    return predictions

  def run(self):
    """Main event loop."""
    print("Action Predictor running...")

    while True:
      # Receive event from logger
      message = self.socket.recv_string()
      event = json.loads(message)

      # Encode and buffer
      state, action = self.encode_event(event)
      self.state_buffer.append(state)
      self.action_buffer.append(action)

      # Predict
      predictions = self.predict_next()

      if predictions:
        print(f"\nPredicted next actions:")
        for p in predictions[:3]:
          print(f"  {p['action']}: {p['probability']:.1%}")


# Usage
if __name__ == '__main__':
  model = load_trained_model('./checkpoints/best')
  predictor = ActionPredictor(model)
  predictor.run()