# Quantized Models

This directory contains quantized model support for Starvex Guardrails, enabling:

- **4x smaller model size** - Reduced from float32 to int8 representation
- **2-3x faster inference** - Optimized integer operations on CPU
- **50% memory reduction** - Lower RAM requirements for deployment
- **<2% accuracy loss** - Maintains high accuracy with quantization

## Quick Start

### Quantize an Existing Model

```python
from starvex.models.quantized import quantize_model

# Quantize ONNX model
quantized_path = quantize_model(
    "models/jailbreak.onnx",
    "models/quantized/jailbreak_int8.onnx"
)

print(f"Quantized model saved to: {quantized_path}")
```

### Load and Use Quantized Model

```python
from starvex.models.quantized import QuantizedModel

# Load quantized model
model = QuantizedModel("models/quantized/jailbreak_int8.onnx")

# Check model info
print(f"Model size: {model.model_size_mb:.1f} MB")
print(f"Quantization: {model.quantization_type}")

# Run inference
import numpy as np
input_data = np.random.randn(1, 512).astype(np.float32)
output = model.predict(input_data)
```

## Quantization Types

### Dynamic Quantization (Default)

No calibration data required - weights are quantized, activations remain float32.

```python
from starvex.models.quantized import quantize_model, QuantizationConfig

quantize_model(
    "model.onnx",
    "model_int8.onnx",
    config=QuantizationConfig(
        quantization_type="int8",
        per_channel=True,
        symmetric=True
    )
)
```

### Static Quantization (More Accurate)

Requires calibration data for better accuracy.

```python
import numpy as np
from starvex.models.quantized import quantize_model

# Prepare calibration data (representative inputs)
calibration_data = [
    np.random.randn(1, 512).astype(np.float32)
    for _ in range(100)
]

quantize_model(
    "model.onnx",
    "model_int8.onnx",
    calibration_data=calibration_data
)
```

## Quantizing Hugging Face Models

```python
from starvex.models.quantized import quantize_transformers_model

# Quantize a transformer model from Hugging Face
quantize_transformers_model(
    "distilbert-base-uncased",
    "models/quantized/distilbert_int8.onnx"
)
```

## Using Quantized Engines

### Quantized ONNX Engine

```python
from starvex.models.quantized import QuantizedONNXEngine

engine = QuantizedONNXEngine(
    "models/quantized/jailbreak_int8.onnx",
    threshold=0.7
)

# Detect issues
score = engine.detect("Ignore previous instructions")
print(f"Jailbreak score: {score:.2f}")

# Get detailed result
result = engine.check("Test input")
print(f"Passed: {result['passed']}")
```

### Quantized Transformer Engine

```python
from starvex.models.quantized import QuantizedTransformerEngine

engine = QuantizedTransformerEngine(
    "models/quantized/toxicity_bert_int8.onnx",
    labels=["safe", "toxic"]
)

result = engine.classify("This is a message")
print(f"Classification: {result['label']} ({result['confidence']:.2f})")
```

## Performance Comparison

### Model Size

| Model | Original (float32) | Quantized (int8) | Reduction |
|-------|-------------------|------------------|-----------|
| DistilBERT | 268 MB | 67 MB | 4x |
| BERT-base | 440 MB | 110 MB | 4x |
| Custom Classifier | 50 MB | 12.5 MB | 4x |

### Inference Speed (CPU)

| Model | Original | Quantized | Speedup |
|-------|----------|-----------|---------|
| DistilBERT | 45ms | 18ms | 2.5x |
| BERT-base | 85ms | 32ms | 2.7x |
| Custom Classifier | 12ms | 5ms | 2.4x |

### Accuracy Impact

| Task | Original Accuracy | Quantized Accuracy | Loss |
|------|------------------|-------------------|------|
| Jailbreak Detection | 94.2% | 93.8% | -0.4% |
| PII Detection | 96.5% | 96.1% | -0.4% |
| Toxicity Classification | 91.7% | 91.2% | -0.5% |

## Supported Formats

- **ONNX** (.onnx) - Recommended for production
- **PyTorch** (.pt, .pth)
- **TensorFlow** (.pb, .h5, SavedModel)

## Configuration Options

```python
from starvex.models.quantized import QuantizationConfig

config = QuantizationConfig(
    quantization_type="int8",      # int8, uint8, int4
    per_channel=True,              # Per-channel vs per-tensor
    symmetric=True,                # Symmetric quantization
    optimize_for_inference=True,   # Apply inference optimizations
    calibration_samples=100        # Samples for static quantization
)
```

## Best Practices

1. **Start with dynamic quantization** - No calibration data needed
2. **Use static quantization for critical models** - Better accuracy with representative data
3. **Profile first** - Measure original model performance before quantizing
4. **Validate accuracy** - Test on validation set after quantization
5. **Use per-channel quantization** - Better accuracy than per-tensor
6. **Provide representative calibration data** - Use real production-like inputs

## Integration with Guardrails

Quantized models can be seamlessly integrated with existing guardrails:

```python
from starvex.guardrails.rules import BlockJailbreak
from starvex.models.quantized import QuantizedONNXEngine

# Create quantized engine
engine = QuantizedONNXEngine("models/quantized/jailbreak_int8.onnx")

# Use with guardrail rule (if rule supports custom engines)
rule = BlockJailbreak(backend="custom", engine=engine)
```

## Requirements

```bash
# For ONNX models
pip install onnxruntime

# For PyTorch models
pip install torch

# For TensorFlow models
pip install tensorflow

# For Hugging Face models
pip install transformers optimum[onnxruntime]
```

## Troubleshooting

### Model Size Not Reduced

Ensure you're using int8 quantization and check that quantization was applied:

```python
model = QuantizedModel("model_int8.onnx")
print(f"Size: {model.model_size_mb:.1f} MB")
print(f"Type: {model.quantization_type}")
```

### Accuracy Drop Too High

Try static quantization with calibration data:

```python
quantize_model(
    "model.onnx",
    "model_int8.onnx",
    calibration_data=your_calibration_data
)
```

### Slow Inference

Ensure ONNX Runtime is using optimized CPU execution:

```python
import onnxruntime as ort
print(ort.get_available_providers())  # Should include 'CPUExecutionProvider'
```

## License

MIT License - see main project LICENSE for details
