Colaboration between @jasonkneen and @ExpressionsBot
Follow us on X
- jasonkneen
- megaprompt
An implementation inspired by Google Research's paper "Generative AI for Programming: A Common Task Framework". This server provides a neural memory system that can learn and predict sequences while maintaining state through a memory vector, following principles outlined in the research for improved code generation and understanding.
? Research Background
This implementation draws from the concepts presented in the Google Research paper (Muennighoff et al., 2024) which introduces a framework for evaluating and improving code generation models. The Titan Memory Server implements key concepts from the paper:
- Memory-augmented sequence learning
- Surprise metric for novelty detection
- Manifold optimization for stable learning
- State maintenance through memory vectors
These features align with the paper's goals of improving code understanding and generation through better memory and state management.
? Features
- Neural memory model with configurable dimensions
- Sequence learning and prediction
- Surprise metric calculation
- Model persistence (save/load)
- Memory state management
- Full MCP tool integration
? Installation
# Install dependencies
npm install
# Build the project
npm run build
# Run tests
npm test
1. ? init_model
Initialize the Titan Memory model with custom configuration.
{
inputDim?: number; // Input dimension (default: 64)
outputDim?: number; // Output/Memory dimension (default: 64)
}
2. ? train_step
Perform a single training step with current and next state vectors.
{
x_t: number[]; // Current state vector
x_next: number[]; // Next state vector
}
3. ? forward_pass
Run a forward pass through the model with an input vector.
{
x: number[]; // Input vector
}
4. ? save_model
Save the model to a specified path.
{
path: string; // Path to save the model
}
5. ? load_model
Load the model from a specified path.
{
path: string; // Path to load the model from
}
6. ℹ️ get_status
Get current model status and configuration.
{} // No parameters required
7. ? train_sequence
Train the model on a sequence of vectors.
{
sequence: number[][]; // Array of vectors to train on
}
? Example Usage
// Initialize model
await callTool('init_model', { inputDim: 64, outputDim: 64 });
// Train on a sequence
const sequence = [
[1, 0, 0, /* ... */],
[0, 1, 0, /* ... */],
[0, 0, 1, /* ... */]
];
await callTool('train_sequence', { sequence });
// Run forward pass
const result = await callTool('forward_pass', {
x: [1, 0, 0, /* ... */]
});
? Technical Details
- Built with TensorFlow.js for efficient tensor operations
- Uses manifold optimization for stable learning
- Implements surprise metric for novelty detection
- Memory management with proper tensor cleanup
- Type-safe implementation with TypeScript
- Comprehensive error handling
? Testing
The project includes comprehensive tests covering:
- Model initialization and configuration
- Training and forward pass operations
- Memory state management
- Model persistence
- Edge cases and error handling
- Tensor cleanup and memory management
Run tests with:
npm test
? Implementation Notes
- All tensor operations are wrapped in
tf.tidy() for proper memory management
- Implements proper error handling with detailed error messages
- Uses type-safe MCP tool definitions
- Maintains memory state between operations
- Handles floating-point precision issues with epsilon tolerance
? License
MIT License - feel free to use and modify as needed!