Metadata-Version: 2.4
Name: chatmix-trainer
Version: 0.1.1
Summary: Train retrieval-based chatbots from structured data
Project-URL: Homepage, https://github.com/opentechcommunity/chatmix-trainer
Project-URL: Documentation, https://github.com/opentechcommunity/chatmix-trainer#readme
Project-URL: Repository, https://github.com/opentechcommunity/chatmix-trainer
Project-URL: Issues, https://github.com/opentechcommunity/chatmix-trainer/issues
Author: ChatMix Contributors
License-Expression: MIT
License-File: LICENSE
Keywords: ai,chatbot,embeddings,retrieval,training
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Requires-Dist: fastembed>=0.3.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: pydantic>=2.0.0
Provides-Extra: dev
Requires-Dist: mypy>=1.0.0; extra == 'dev'
Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
Requires-Dist: pytest>=7.0.0; extra == 'dev'
Requires-Dist: ruff>=0.1.0; extra == 'dev'
Requires-Dist: twine>=4.0.0; extra == 'dev'
Provides-Extra: faiss
Requires-Dist: faiss-cpu>=1.7.4; extra == 'faiss'
Description-Content-Type: text/markdown

# chatmix-trainer

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

[![PyPI version](https://badge.fury.io/py/chatmix-trainer.svg)](https://badge.fury.io/py/chatmix-trainer)
[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/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](https://github.com/qdrant/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

```bash
pip install chatmix-trainer
```

### Basic Example

```python
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

```python
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

```python
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

```python
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

```python
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

```python
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)
```

---

```python
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

```python
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`:

```python
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:

```python
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):

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

### Command-Line Interface

Train models from the terminal:

```bash
# 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

```python
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](https://qdrant.github.io/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](https://qdrant.github.io/fastembed/examples/Supported_Models/).

**Change model:**

```python
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:**

```python
# 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**:

```python
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

```python
# 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:

```bash
pip install chatmix-trainer
```

### Optional Dependencies

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

```bash
pip install chatmix-trainer[faiss]
```

---

```bash
# 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

```bash
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

```bash
# 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](https://github.com/opentechcommunity/chatmix-trainer/blob/master/LICENSE) for details.

Free for personal and commercial use.

---

## Links

- **PyPI**: https://pypi.org/project/chatmix-trainer/
- **GitHub**: https://github.com/opentechcommunity/chatmix-trainer
- **Issues**: https://github.com/opentechcommunity/chatmix-trainer/issues
- **Discussions**: https://github.com/opentechcommunity/chatmix-trainer/discussions

### Related Projects

- **[chatmix-widget](https://www.npmjs.com/package/chatmix-widget)** - Embeddable chat UI widget (npm)
- **[FastEmbed](https://github.com/qdrant/fastembed)** - Fast CPU embeddings library
- **[Qdrant](https://qdrant.tech/)** - Vector database (for production scale)

---

## Show Your Support

If you find chatmix-trainer useful:

- Star the [GitHub repository](https://github.com/opentechcommunity/chatmix-trainer)
- Report bugs and request features via [Issues](https://github.com/opentechcommunity/chatmix-trainer/issues)
- Join the [Discussions](https://github.com/opentechcommunity/chatmix-trainer/discussions)
- Share your project built with chatmix-trainer!

---

## Stats

![PyPI Downloads](https://img.shields.io/pypi/dm/chatmix-trainer)
![GitHub Stars](https://img.shields.io/github/stars/opentechcommunity/chatmix-trainer?style=social)
![GitHub Issues](https://img.shields.io/github/issues/opentechcommunity/chatmix-trainer)

---

**Built with ❤️ by the ChatMix team**

[Website](https://github.com/opentechcommunity/chatmix-trainer) • [Documentation](https://github.com/opentechcommunity/chatmix-trainer)
