SAGE Dialogue Gen 🌱

Authors: Yizhe Zhang, Navdeep Jaitly (Apple)

Model Information

Language: English
License: Apache 2.0
Base Model: mistralai/Mixtral-8x7B-Instruct-v0.1
Library: transformers
Tags: dialog-generation, conversational-ai, state-action-model
Dataset: ShareGPT
Metrics: Custom emotional-intelligence evaluation

Citation

@misc{zhang2025sage,
  title = {SAGE: Steering and Refining Dialogue Generation with State‑Action Augmentation},
  author = {Zhang, Yizhe and Jaitly, Navdeep},
  year = {2025},
  howpublished = {arXiv preprint},
  note = {arXiv:2503.03040}
}

📄 Paper: Available on arXiv and Papers with Code

Model Description

SAGE introduces latent state-action variables between dialogue turns, enabling:

Structured Control: Precise management of emotional tone and conversational strategy
Enhanced Emotional Intelligence: Explicit state planning for more empathetic responses
Self-Improving Pipeline: Comprehensive training approach including:
- Data augmentation
- Dialogue-tree search
- Reward modeling
- Fine-tuning optimization

This approach allows for more nuanced and contextually appropriate dialogue generation compared to traditional methods.

Intended Uses

✅ Recommended Applications

Emotional or empathetic chatbots
Long-horizon, strategy-aware conversation systems
Research on structured latent-variable dialogue control
Educational conversational AI systems
Customer service applications requiring emotional intelligence

⚠️ Important Limitations

Not suitable for high-stakes, safety-critical deployment without further evaluation
Requires additional testing for production environments
May need domain-specific fine-tuning for specialized applications

Training Details

Base Model: Mixtral-8x7B-Instruct

Training Pipeline:

Data Preparation: ShareGPT-style JSON formatting
Supervised Fine-Tuning (SFT): Initial model adaptation
Dialogue-Tree Search: Exploration of conversation paths
Preference Learning: Reward model training
Comparative Evaluation: Performance assessment and inference optimization

Performance

SAGE demonstrates significant improvements on emotional-intelligence metrics compared to baseline models while maintaining generative flexibility and coherence. The model shows particular strength in:

Emotional tone consistency
Contextual appropriateness
Long-term conversation planning
Empathetic response generation

Usage

Quick Start

git clone https://github.com/apple/ml-sage-dialog-gen
cd ml-sage-dialog-gen
bash setup.sh

Basic Implementation

from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model
tokenizer = AutoTokenizer.from_pretrained("apple/sage-dialogue-gen")
model = AutoModelForCausalLM.from_pretrained("apple/sage-dialogue-gen")

# Generate dialogue
input_text = "I'm feeling overwhelmed with work lately."
inputs = tokenizer(input_text, return_tensors="pt")
outputs = model.generate(**inputs, max_length=150, do_sample=True, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)

Requirements

Python 3.8+
PyTorch 1.12+
Transformers 4.21+
Additional dependencies listed in requirements.txt

Contributing

Contributions are welcome! Please see our contributing guidelines and code of conduct before submitting pull requests.

License

This project is licensed under the Apache License 2.0. See the LICENSE file for details.

Acknowledgments

Built upon the Mixtral-8x7B-Instruct foundation model
Trained using ShareGPT dataset
Developed by the Apple Machine Learning Research team

Contact

For questions or issues, please open a GitHub issue or contact the development team through the official Apple ML research channels.

apple
/

sage-ft-mixtral-8x7b