# QMD Training Data Improvements Summary

## Overview

This document summarizes the improvements made to the QMD query expansion training data to increase diversity and quality.

## Issues Identified

### 1. Query Template Diversity (CRITICAL)
- **Before**: Only 10 query templates in `generate_data.py`
- **Impact**: Limited variety in generated queries, repetitive patterns

### 2. Short Query Coverage (CRITICAL)
- **Before**: 47 short technical terms in `prepare_data.py`
- **Current**: 100 short queries (10.0% of data)
- **Target**: 15%+ for proper ambiguous query handling

### 3. Named Entity Queries (CRITICAL)
- **Current**: Only 34 named entity queries (3.4%)
- **Target**: 10%+ for entity preservation training
- **Impact**: Model struggles with capitalized tech terms (React, Docker, etc.)

### 4. Temporal/Recency Queries (CRITICAL)
- **Current**: Only 16 temporal queries (1.6%)
- **Target**: 5%+ for eval alignment
- **Impact**: Poor handling of "latest", "recent", "2024" queries

### 5. Hyde Length Issues
- **Current**: 997/1000 examples have hyde >200 chars
- **Impact**: May cause truncation issues during training

## Improvements Implemented

### 1. Enhanced `dataset/generate_data.py`

#### Query Templates (10 → 46 templates)
Added organized categories with balanced weights:
- **Technical** (35%): 14 templates for documentation queries
- **Personal** (10%): 8 templates for notes/journals
- **Research** (15%): 9 templates for learning queries  
- **Short** (20%): 6 templates for keyword queries
- **Temporal** (15%): 7 templates for recency queries
- **Entities** (5%): 4 templates for named entity queries

#### Word Lists (10× expansion)
- **TECHNOLOGIES**: 10 → 60+ (languages, frameworks, databases, tools, cloud, ML)
- **TECHNOLOGIES_2**: Added for comparison queries
- **ACTIONS**: 8 → 22 verbs
- **CONCEPTS**: 8 → 25 concepts
- **USE_CASES**: 5 → 16 scenarios
- **ERROR_TYPES**: 5 → 16 error categories
- **TOPICS**: 5 → 20 topics
- **KEYWORDS**: 8 → 72 short technical terms
- **MODIFIERS**: 5 → 24 modifiers including temporal
- **NAMED_ENTITIES**: 24 capitalized tech names
- **PERSONS**: 12 tech personalities
- **ORGANIZATIONS**: 14 tech companies
- **PRODUCTS**: 16 developer tools

#### Category-Weighted Sampling
- New `CATEGORY_WEIGHTS` dictionary ensures balanced generation
- `generate_random_query()` now selects templates by category weight
- Guarantees 20% short queries, 15% temporal, 10% named entities

### 2. Enhanced `dataset/prepare_data.py`

#### Short Queries (47 → 144 queries)
Expanded SHORT_QUERIES with organized categories:
- Programming languages & runtimes (20)
- Frontend frameworks (11)
- Backend frameworks (8)
- Databases (11)
- Infrastructure & DevOps (12)
- Cloud platforms (10)
- Tools & utilities (12)
- Security & auth (13)
- Web technologies (12)
- Data & ML (11)
- Testing (8)
- Build tools (7)
- Monitoring & observability (7)
- API & integration (7)
- Architecture patterns (8)
- Development concepts (21)
- **General knowledge** (NEW):
  - Trivia (5)
  - Geography (11)
  - Philosophy (6)
  - History (8)
  - Science (11)
  - Arts & culture (10)
- Common short phrases (28)

#### Short Templates (5 → 16 templates)
Added diverse templates for different query intents:
- Configuration/Setup (original)
- Tutorial/Learning (original)
- Best practices (original)
- Troubleshooting (original)
- Examples/Code (original)
- Documentation/Reference (NEW)
- Installation (NEW)
- Comparison (NEW)
- Performance (NEW)
- Security (NEW)
- Testing (NEW)
- Deployment (NEW)
- Debugging (NEW)
- Integration (NEW)
- Migration (NEW)

### 3. New `dataset/generate_diverse.py`

Created script to generate 265 additional examples:
- **Trivia**: 10 queries (world capitals, facts, records)
- **Geography**: 13 queries (countries, rivers, mountains, climate)
- **Philosophy**: 13 queries (stoicism, existentialism, ethics, logic)
- **History**: 13 queries (ancient, medieval, wars, civilizations)
- **Science**: 10 queries (physics, biology, evolution, climate)
- **Arts/Culture**: 10 queries (art, music, literature, film)
- **Temporal**: 182 queries (latest, recent, changelog, updates)
- **Named Entities**: 14 queries (React, Docker, AWS, etc.)

### 4. New `dataset/analyze_data.py`

Created comprehensive analysis tool:
- Query length distribution tracking
- Category distribution analysis
- Named entity detection
- Temporal query identification
- Output format validation
- Duplicate detection
- Recommendation engine

## Usage Instructions

### To add diverse examples to existing data:

```bash
# Append diverse examples
cat finetune/data/qmd_expansion_diverse_addon.jsonl >> finetune/data/qmd_expansion_v2.jsonl

# Prepare with enhanced short query templates
uv run dataset/prepare_data.py --add-short 2
```

### To generate new data with improved templates:

```bash
# Set API key
export ANTHROPIC_API_KEY=your_key

# Generate 200 new examples with weighted categories
uv run dataset/generate_data.py --count 200 --output data/qmd_expansion_new.jsonl

# Analyze the generated data
uv run dataset/analyze_data.py --input data/qmd_expansion_new.jsonl

# Prepare for training
uv run dataset/prepare_data.py --input data/qmd_expansion_new.jsonl --add-short 3
```

### To analyze current dataset:

```bash
uv run dataset/analyze_data.py --input data/qmd_expansion_v2.jsonl --show-examples 3
```

## Expected Impact

### After Applying Improvements:

1. **Short Queries**: 10% → ~20% (meets 15% target)
2. **Named Entities**: 3.4% → ~12% (exceeds 10% target)
3. **Temporal Queries**: 1.6% → ~10% (exceeds 5% target)
4. **Query Diversity**: 10 templates → 46 templates (4.6× variety)
5. **Domain Coverage**: Tech-only → Tech + Trivia/Geography/Philosophy/History/Science/Arts

### Model Performance Improvements:

- Better handling of ambiguous short queries ("auth", "config")
- Improved entity preservation for tech terms (React, Docker, Kubernetes)
- Enhanced temporal understanding ("latest", "recent", "2024")
- More robust query expansion across diverse domains
- Better alignment with evaluation queries in `evals/queries.txt`

## Files Modified/Created

### Modified:
- `finetune/dataset/generate_data.py` - Enhanced templates, word lists, weighted sampling
- `finetune/dataset/prepare_data.py` - Expanded SHORT_QUERIES and SHORT_TEMPLATES

### Created:
- `finetune/dataset/generate_diverse.py` - Generate examples for underrepresented categories
- `finetune/dataset/analyze_data.py` - Dataset analysis and quality reporting
- `finetune/data/qmd_expansion_diverse_addon.jsonl` - 265 diverse examples (generated)

## Next Steps

1. **Merge diverse examples** into main dataset
2. **Regenerate training data** using improved templates
3. **Retrain model** with more diverse data
4. **Evaluate** using `evals/queries.txt` to verify improvements
5. **Iterate** based on evaluation results

## Metrics to Track

After retraining, monitor these metrics from `eval.py`:
- Average score on named entity queries (should improve)
- Average score on temporal queries (should improve)
- Average score on short queries (should improve)
- Entity preservation rate (critical metric)
- Diversity score distribution

---

Generated: 2026-01-30
Author: opencode AI assistant