Skip to main content

Train retrieval-based chatbots from structured data

Project description

chatmix-trainer

Train retrieval-based chatbots from structured Q&A data — no GPU required

PyPI version Python 3.9+ License: MIT

chatmix-trainer is a Python library for building semantic search-powered chatbots from structured Q&A data. It uses CPU-optimized embeddings (ONNX via FastEmbed) to create fast, portable chatbot models that work anywhere—no GPU, no cloud, no dependencies on LLM APIs.

Perfect for:

  • FAQ chatbots
  • Knowledge base search
  • Customer support automation
  • Documentation helpers
  • Internal company wikis
  • Any retrieval-based Q&A system

Features

  • Fast CPU Training — Train on your laptop in seconds, no GPU required
  • Portable Artifacts — Single .artifact file, deploy anywhere
  • Semantic Search — Understanding-based matching, not just keywords
  • High Precision — Returns only relevant results above confidence threshold
  • Minimal Dependencies — No PyTorch, TensorFlow, or heavy ML frameworks
  • Instant Reload — Load pre-trained models instantly without retraining
  • Framework Agnostic — Use with FastAPI, Flask, Django, or standalone
  • Type Safe — Full type hints and Pydantic validation
  • Well Tested — Comprehensive test suite included

Quick Start

Installation

pip install chatmix-trainer

Basic Example

from chatmix_trainer import ChatbotModel, Trainer, TrainConfig

# 1. Prepare your Q&A data
data = [
    {
        "question": "What are your business hours?",
        "answer": "We're open Monday-Friday, 9 AM to 5 PM EST."
    },
    {
        "question": "How do I reset my password?",
        "answer": "Click 'Forgot Password' on the login page and follow the email instructions."
    },
    {
        "question": "Do you offer refunds?",
        "answer": "Yes, we offer full refunds within 30 days of purchase."
    }
]

# 2. Train the model (takes ~5 seconds)
trainer = Trainer(TrainConfig())
model = trainer.train(data, output_path="chatbot.artifact")

# 3. Query the model
results = model.query("when are you open?", top_k=1)
print(results[0]["answer"])
# Output: "We're open Monday-Friday, 9 AM to 5 PM EST."

That's it! You now have a working chatbot that understands semantic meaning.


How It Works

User Question
     ↓
[Embedding Model]  ← Converts text to vectors
     ↓
[Vector Search]    ← Finds similar Q&A pairs
     ↓
[Ranked Results]   ← Returns best matches
     ↓
Answer(s)

chatmix-trainer uses semantic embeddings to understand the meaning of questions, not just keywords. This means:

  • "when are you open?" matches "What are your business hours?"
  • "how do I change my password?" matches "How do I reset my password?"
  • Works even with typos and different phrasings

Usage Examples

1. Command-Line Chatbot

from chatmix_trainer import ChatbotModel

model = ChatbotModel.load("chatbot.artifact")

print("Chatbot ready! Type 'quit' to exit.\n")
while True:
    question = input("You: ")
    if question.lower() in ['quit', 'exit']:
        break
    
    results = model.query(question, top_k=1)
    if results:
        print(f"Bot: {results[0]['answer']}\n")
    else:
        print("Bot: I don't have an answer for that.\n")

2. Web API with FastAPI

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from chatmix_trainer import ChatbotModel
import json

app = FastAPI()
model = ChatbotModel.load("chatbot.artifact")

@app.post("/chat")
async def chat(request: dict):
    question = request["message"]
    results = model.query(question, top_k=1)
    
    if results:
        answer = results[0]["answer"]
    else:
        answer = "I don't have an answer for that."
    
    # Stream response token by token
    async def generate():
        for word in answer.split():
            yield f"data: {json.dumps({'token': word + ' '})}\n\n"
        yield f"data: {json.dumps({'done': True})}\n\n"
    
    return StreamingResponse(generate(), media_type="text/event-stream")

# Run with: uvicorn server:app --reload

3. Slack Bot

from slack_bolt import App
from chatmix_trainer import ChatbotModel

app = App(token="xoxb-your-token")
model = ChatbotModel.load("chatbot.artifact")

@app.message(".*")
def handle_message(message, say):
    question = message['text']
    results = model.query(question, top_k=1)
    
    if results:
        say(results[0]["answer"])
    else:
        say("I don't have an answer for that. Try rephrasing?")

if __name__ == "__main__":
    app.start(port=3000)

4. Discord Bot

import discord
from chatmix_trainer import ChatbotModel

client = discord.Client(intents=discord.Intents.default())
model = ChatbotModel.load("chatbot.artifact")

@client.event
async def on_message(message):
    if message.author == client.user:
        return
    
    results = model.query(message.content, top_k=1)
    if results:
        await message.channel.send(results[0]["answer"])

client.run("YOUR_BOT_TOKEN")

5. Flask Integration

from flask import Flask, request, jsonify
from chatmix_trainer import ChatbotModel

app = Flask(__name__)
model = ChatbotModel.load("chatbot.artifact")

@app.route('/ask', methods=['POST'])
def ask():
    question = request.json.get('question', '')
    results = model.query(question, top_k=3)
    
    return jsonify({
        'results': [
            {
                'answer': r['answer'],
                'score': r['score'],
                'question': r.get('question', '')
            }
            for r in results
        ]
    })

if __name__ == '__main__':
    app.run(debug=True)

from chatmix_trainer import Trainer, TrainConfig

# Configure training (all optional)
config = TrainConfig(
    embedding_model="BAAI/bge-small-en-v1.5",  # FastEmbed model name
    chunk_size=512,                             # Max characters per chunk
    top_k=3,                                    # Default results to return
    similarity_threshold=0.3,                   # Min score to return (0.0-1.0)
)

# Train
trainer = Trainer(config)
result = trainer.train(
    data=[
        {"id": "1", "question": "...", "answer": "..."},
        {"id": "2", "question": "...", "answer": "..."},
    ],
    output_path="bot.artifact",
    on_complete=lambda r: print(f"Done! Saved to {r.artifact_path}"),
    on_progress=lambda done, total: print(f"{done}/{total}"),
)

print(f"Trained {result.record_count} records in {result.elapsed_seconds:.2f}s")
print(f"Checksum: {result.checksum}")

Retrieval

from chatmix_trainer import ChatbotModel

# Load trained model
model = ChatbotModel.from_file("bot.artifact")

# Retrieve relevant records
matches = model.retrieve(
    query="what are your hours?",
    k=5,              # Override default top_k
    threshold=0.4,    # Override default threshold
)

# Each match contains:
for match in matches:
    print(match.record.id)         # Record ID
    print(match.record.question)   # Question text
    print(match.record.answer)     # Answer text
    print(match.record.tags)       # Tags list
    print(match.record.metadata)   # Metadata dict
    print(match.score)             # Similarity score (0.0-1.0)
    print(match.chunk_index)       # Which chunk matched (if chunked)

Advanced Usage

Training Configuration

Customize training behavior with TrainConfig:

from chatmix_trainer import Trainer, TrainConfig

config = TrainConfig(
    embedding_model="BAAI/bge-small-en-v1.5",  # Model to use
    chunk_size=512,                             # Max chars per chunk
    chunk_overlap=50,                           # Overlap between chunks
    similarity_threshold=0.3,                   # Min score (0.0-1.0)
    top_k=3,                                    # Default results
    normalize_embeddings=True,                  # L2 normalization
    case_sensitive=False,                       # Case matching
)

trainer = Trainer(config)
model = trainer.train(data, output_path="bot.artifact")

Data Format

Your training data should be a list of dictionaries with these fields:

data = [
    {
        "id": "unique-id-1",              # Required: unique identifier
        "question": "What is Python?",    # Required: the question
        "answer": "Python is...",         # Required: the answer
        "tags": ["programming", "python"], # Optional: for filtering
        "metadata": {                     # Optional: custom data
            "category": "programming",
            "difficulty": "beginner"
        }
    }
]

Alternative format (single text field):

data = [
    {
        "id": "doc-1",
        "text": "Python is a high-level programming language...",
        "tags": ["python", "intro"]
    }
]

Command-Line Interface

Train models from the terminal:

# Basic training
chatmix-train --input data.json --output chatbot.artifact

# With options
chatmix-train \
  --input data.json \
  --output chatbot.artifact \
  --model "BAAI/bge-base-en-v1.5" \
  --threshold 0.4 \
  --top-k 5 \
  --verbose

# Get help
chatmix-train --help

Querying

from chatmix_trainer import ChatbotModel

model = ChatbotModel.load("chatbot.artifact")

# Basic query
results = model.query("How do I reset my password?")

# With options
results = model.query(
    query="How do I reset my password?",
    top_k=5,                    # Return top 5 matches
    threshold=0.4,              # Only results above 0.4 score
    filter_tags=["account"],    # Filter by tags
)

# Access results
for result in results:
    print(f"Question: {result['question']}")
    print(f"Answer: {result['answer']}")
    print(f"Score: {result['score']:.3f}")
    print(f"Tags: {result.get('tags', [])}")
    print(f"Metadata: {result.get('metadata', {})}")
    print()

Embedding Models

chatmix-trainer uses FastEmbed for CPU-optimized embeddings. Popular models:

Model Size Speed Quality Best For
BAAI/bge-small-en-v1.5 133MB Fast Good Default - English
sentence-transformers/all-MiniLM-L6-v2 90MB Very Fast Decent Speed-critical apps
BAAI/bge-base-en-v1.5 436MB Medium Better Higher quality needed
BAAI/bge-large-en-v1.5 1.3GB Slow Best Maximum accuracy

For multilingual support, see FastEmbed's model list.

Change model:

config = TrainConfig(embedding_model="BAAI/bge-base-en-v1.5")

The Artifact File

The .artifact file is a portable ZIP archive containing:

chatbot.artifact/
  ├── embeddings.npz      # Compressed embedding matrix
  ├── records.json        # Original Q&A data
  ├── record_map.json     # Index mapping
  ├── config.json         # Training configuration
  └── metadata.json       # Version, checksums, stats

Benefits:

  • Single file — easy to version control and deploy
  • Portable — works on any machine with Python 3.9+
  • Fast loading — instant startup, no retraining
  • Reproducible — same artifact = same results
  • Small size — typically <10MB for hundreds of Q&A pairs

Usage pattern:

# Train once (development/CI)
trainer = Trainer()
model = trainer.train(data, output_path="v1.artifact")

# Deploy anywhere (production)
model = ChatbotModel.load("v1.artifact")  # Instant!
results = model.query("...")

Integration with LLMs (RAG Pattern)

chatmix-trainer handles retrieval — pair it with an LLM for generation:

from chatmix_trainer import ChatbotModel
import openai  # or use Ollama, Anthropic, etc.

model = ChatbotModel.load("chatbot.artifact")

def answer_with_llm(question: str) -> str:
    # 1. Retrieve relevant context
    results = model.query(question, top_k=3)
    
    # 2. Format context for LLM
    context = "\n\n".join([
        f"Q: {r['question']}\nA: {r['answer']}"
        for r in results
    ])
    
    # 3. Generate natural answer with LLM
    prompt = f"""Use this context to answer the question.
    
Context:
{context}

Question: {question}

Answer:"""
    
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}]
    )
    
    return response.choices[0].message.content

# Use it
answer = answer_with_llm("What are your business hours?")
print(answer)

This RAG (Retrieval-Augmented Generation) pattern:

  • Reduces LLM hallucinations
  • Grounds answers in your data
  • Works with any LLM (OpenAI, Claude, Ollama, etc.)
  • More cost-effective than fine-tuning

Use Cases

Perfect For

  • FAQ Chatbots — Customer support, product questions
  • Documentation Search — Help users find docs faster
  • Internal Knowledge Base — Company wikis, procedures
  • Customer Support — Automate common questions
  • E-commerce — Product questions, shipping, returns
  • Education — Course FAQs, tutoring bots
  • Healthcare — Patient information, appointment scheduling
  • IT Helpdesk — Password resets, troubleshooting

Real-World Examples

# E-commerce Support
data = [
    {"question": "How do I track my order?", "answer": "..."},
    {"question": "What's your return policy?", "answer": "..."},
    {"question": "Do you ship internationally?", "answer": "..."},
]

# IT Helpdesk
data = [
    {"question": "How do I reset my password?", "answer": "..."},
    {"question": "VPN not connecting", "answer": "..."},
    {"question": "Install Microsoft Office", "answer": "..."},
]

# Restaurant FAQ
data = [
    {"question": "What are your hours?", "answer": "..."},
    {"question": "Do you take reservations?", "answer": "..."},
    {"question": "Is there parking?", "answer": "..."},
]

Performance

Training Speed

Dataset Size Time Memory
100 Q&A pairs ~2 seconds ~200MB
1,000 Q&A pairs ~15 seconds ~500MB
10,000 Q&A pairs ~2 minutes ~2GB

On Apple M1 Pro, Python 3.11

Query Speed

  • Without LLM: 10-50ms per query
  • With LLM: 2-10 seconds (depends on LLM)

Artifact Size

  • ~100KB per 100 Q&A pairs
  • Compressed with gzip
  • Typically <10MB for most use cases

Requirements

  • Python: 3.9 or higher
  • CPU: Any modern CPU (no GPU needed)
  • Memory: ~500MB minimum, ~2GB for large datasets
  • Storage: ~10MB per 1000 Q&A pairs

Dependencies

  • fastembed - CPU-optimized embeddings
  • numpy - Array operations
  • pydantic - Data validation

All dependencies are automatically installed with:

pip install chatmix-trainer

Optional Dependencies

For very large datasets (100k+ records), install FAISS for faster search:

pip install chatmix-trainer[faiss]

# Clone the repo
git clone https://github.com/opentechcommunity/chatmix-trainer.git
cd chatmix-trainer

# Create virtual environment with uv
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install in editable mode with dev dependencies
uv pip install -e ".[dev]"

# Run tests
pytest

# Lint
ruff check .

# Format
ruff format .

# Type check
mypy src

Testing

pytest                    # Run all tests
pytest --cov             # With coverage
pytest -v                # Verbose
pytest tests/test_model.py  # Single file

FAQ

Q: Do I need a GPU?
A: No. All embeddings run on CPU using ONNX optimization.

Q: Can I use this without an LLM?
A: Yes! The best matching answer often works as-is. LLMs are optional for more natural responses.

Q: How big can my dataset be?
A: Tested with 10k+ Q&A pairs. For 100k+, install chatmix-trainer[faiss] for faster search.

Q: Does this fine-tune a language model?
A: No. This is retrieval-based, not fine-tuning. It builds a searchable index of your data.

Q: Can I update the training data?
A: Yes — retrain with new data and redeploy the artifact. Training is fast enough to run on-demand.

Q: What if my Q&A pairs are very long?
A: Set chunk_size in TrainConfig to split them into smaller pieces. Each chunk is embedded separately.

Q: How accurate is it?
A: Depends on your data quality and the embedding model. Generally very good for domain-specific Q&A.

Q: Can I customize the similarity threshold?
A: Yes, both in TrainConfig (default) and at query time via the threshold parameter.

Q: Does it work offline?
A: Yes! After initial model download, everything runs locally. No API calls required.

Q: Can I use it commercially?
A: Yes! MIT license allows commercial use.


🤝 Contributing

Contributions are welcome! Please open an issue or pull request on GitHub.

Development Setup

# Clone repository
git clone https://github.com/opentechcommunity/chatmix-trainer.git
cd chatmix-trainer

# Create virtual environment
python -m venv venv
source venv/bin/activate  # Windows: venv\Scripts\activate

# Install in editable mode with dev dependencies
pip install -e ".[dev]"

# Run tests
pytest

# Run type checking
mypy src

# Run linter
ruff check .

# Format code
ruff format .

License

MIT License - see LICENSE for details.

Free for personal and commercial use.


Links

Related Projects


Show Your Support

If you find chatmix-trainer useful:


Stats

PyPI Downloads GitHub Stars GitHub Issues


Built with ❤️ by the ChatMix team

WebsiteDocumentation

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

chatmix_trainer-0.1.1.tar.gz (16.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

chatmix_trainer-0.1.1-py3-none-any.whl (17.2 kB view details)

Uploaded Python 3

File details

Details for the file chatmix_trainer-0.1.1.tar.gz.

File metadata

  • Download URL: chatmix_trainer-0.1.1.tar.gz
  • Upload date:
  • Size: 16.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.3

File hashes

Hashes for chatmix_trainer-0.1.1.tar.gz
Algorithm Hash digest
SHA256 86f93088de4734d8965bcf76345c4572071356f39d2648ae987af4f4ae67f391
MD5 200fe183c55fccebb69ac21fea648a02
BLAKE2b-256 d63ad9906b03953594980c5a3b438a1dc9a3f77b5b4d1afcf9fad28a98312a0b

See more details on using hashes here.

File details

Details for the file chatmix_trainer-0.1.1-py3-none-any.whl.

File metadata

File hashes

Hashes for chatmix_trainer-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 2773cec7b3d5a0b7a403a6fd417164f5ba6fa0e99e23cdd7373b467ad0e5b6d1
MD5 d4a7d153450f7028284e6c581b0d6714
BLAKE2b-256 6f82f5e9c8c46a1ae343315b89520e4cffd1de126c8a6ce2747785b129dd5877

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page