BitNet b1.58 2B4T - Rust Implementation

A native Rust implementation of BitNet b1.58 2B4T model for efficient 1-bit LLM inference using the Candle framework.

Features

• 🚀 Native Rust implementation with Candle
• 🍎 Apple Silicon (Metal) GPU acceleration support
• 💬 Multiple interaction modes (CLI, interactive, programmatic)
• 🔧 Simple API for integration
• ⚡ Efficient 1-bit quantized inference

Installation

git clone <repository>
cd bitnet-rust
cargo build --release

Usage

Command Line Interface

Single Query

# Ask a single question
cargo run "What is the capital of France?"

# Or using the binary name
cargo run --bin bitnet "Explain quantum computing"

Interactive Mode

# Start interactive chat session
cargo run -- --interactive
# or
cargo run -- -i

Prompt for Input

# Run without arguments to be prompted for input
cargo run

Programmatic Usage

use bitnet_rust::{BitNetGenerator, BitNetConfig, BitNetForCausalLM};
use anyhow::Result;

fn main() -> Result<()> {
    // Load model (see examples/simple_chat.rs for full setup)
    let generator = BitNetGenerator::new(model, tokenizer, device);
    
    // Generate response
    let response = generator.generate("What is machine learning?")?;
    println!("Response: {}", response);
    
    // Generate with custom system prompt
    let response = generator.generate_with_system(
        "Write a poem",
        Some("You are a creative poet")
    )?;
    println!("Poem: {}", response);
    
    Ok(())
}

Examples

Run the included example:

cargo run --example simple_chat

API Reference

BitNetGenerator

The main interface for text generation:

impl BitNetGenerator {
    // Create new generator
    pub fn new(model: BitNetForCausalLM, tokenizer: Tokenizer, device: Device) -> Self;
    
    // Generate response with default system prompt
    pub fn generate(&self, message: &str) -> Result<String>;
    
    // Generate response with custom system prompt
    pub fn generate_with_system(&self, message: &str, system: Option<&str>) -> Result<String>;
    
    // Get the device being used
    pub fn device(&self) -> &Device;
}

Utility Functions

// Format Llama 3 chat template
pub fn format_llama3_chat(system: &str, user_message: &str) -> String;

// Interactive chat session
pub fn run_interactive_chat<F>(generate_fn: F) -> Result<()>
where F: FnMut(&str) -> Result<String>;

ConversationHistory

For multi-turn conversations:

let mut conv = ConversationHistory::new("You are helpful".to_string());
conv.add_user_message("Hello".to_string());
conv.add_assistant_message("Hi there!".to_string());
let prompt = conv.get_prompt();

Model Configuration

The model automatically downloads from Hugging Face Hub:

• Model: microsoft/bitnet-b1.58-2B-4T
• Architecture: BitNet with 1.58-bit quantization
• Parameters: 2B parameters
• Context length: 4096 tokens

Performance

• Apple Silicon: Optimized for M1/M2/M3 chips with Metal acceleration
• CPU: Fallback support for other platforms
• Memory: Efficient 1-bit quantization reduces memory usage significantly

Examples Output

🚀 BitNet b1.58 2B4T - Native 1-bit LLM
============================================================
✓ Using Apple GPU (Metal) - optimal for M1/M2/M3 Macs

💬 Query: What is the capital of France?

⚡ Generating response...

============================================================
🤖 Response: The capital of France is Paris. It is located in the north-central part of the country and is the largest city in France, serving as the political, economic, and cultural center of the nation.
============================================================
⏱️  Time: 2.34s | Device: Metal(0)

Dependencies

• candle-core & candle-nn: ML framework
• hf-hub: Hugging Face model loading
• tokenizers: Text tokenization
• anyhow: Error handling
• serde: JSON parsing
• rand: Random sampling

License

[Add your license here]

Contributing

[Add contribution guidelines here]