qmdr 1.0.0

package/finetune/dataset/generate_ollama.py

@@ -0,0 +1,326 @@
+#!/usr/bin/env python3
+"""Generate synthetic training data for QMD query expansion using local Ollama."""
+
+import argparse
+import json
+import random
+import sys
+import time
+
+from dataset.schema import normalize_output_items, parse_output_text
+from pathlib import Path
+
+try:
+    import requests
+except ImportError:
+    print("Install requests: pip install requests")
+    exit(1)
+
+# Diverse query seeds across many domains
+QUERY_SEEDS = [
+    # Programming & Tech
+    "async await javascript",
+    "rust ownership borrow checker",
+    "kubernetes pod networking",
+    "docker compose volumes",
+    "nginx reverse proxy",
+    "postgresql index optimization",
+    "redis caching strategies",
+    "graphql mutations",
+    "websocket authentication",
+    "terraform state management",
+    "ansible playbook variables",
+    "prometheus alerting rules",
+    "elasticsearch aggregations",
+    "kafka consumer groups",
+    "grpc streaming",
+    "oauth2 refresh tokens",
+    "jwt token expiration",
+    "cors preflight requests",
+    "css grid layout",
+    "react hooks useEffect",
+    "vue composition api",
+    "svelte stores",
+    "nextjs middleware",
+    "webpack code splitting",
+    "typescript generics constraints",
+    "python asyncio gather",
+    "go goroutines channels",
+    "java streams filter map",
+    "c++ smart pointers",
+    "swift optionals unwrapping",
+    # DevOps & Infrastructure
+    "ci cd pipeline best practices",
+    "blue green deployment",
+    "canary release strategy",
+    "infrastructure as code",
+    "secrets management vault",
+    "load balancer health checks",
+    "ssl certificate renewal",
+    "dns propagation time",
+    "cdn cache invalidation",
+    "container orchestration",
+    "service mesh istio",
+    "observability tracing",
+    "log aggregation elk",
+    "metrics dashboards grafana",
+    "incident response runbook",
+    # Data & ML
+    "pandas dataframe groupby",
+    "numpy array broadcasting",
+    "scikit learn pipeline",
+    "pytorch autograd",
+    "tensorflow keras layers",
+    "huggingface transformers",
+    "feature engineering techniques",
+    "hyperparameter tuning",
+    "model evaluation metrics",
+    "data preprocessing normalization",
+    "time series forecasting",
+    "anomaly detection",
+    "recommendation systems",
+    "natural language processing",
+    "computer vision cnn",
+    "reinforcement learning",
+    "transfer learning",
+    "model deployment mlops",
+    # Databases
+    "sql join types explained",
+    "database normalization forms",
+    "acid transactions",
+    "database sharding strategies",
+    "read replicas setup",
+    "connection pooling",
+    "query optimization explain",
+    "stored procedures triggers",
+    "database migrations",
+    "nosql document model",
+    "graph database queries",
+    "vector database similarity",
+    # Security
+    "xss prevention sanitization",
+    "sql injection prepared statements",
+    "csrf tokens",
+    "content security policy",
+    "rate limiting api",
+    "input validation patterns",
+    "password hashing bcrypt",
+    "two factor authentication",
+    "penetration testing",
+    "security headers http",
+    "vulnerability scanning",
+    "audit logging",
+    # System Administration
+    "linux file permissions",
+    "systemd service unit",
+    "cron job scheduling",
+    "ssh key management",
+    "firewall rules iptables",
+    "process monitoring",
+    "disk space management",
+    "memory leak debugging",
+    "network troubleshooting",
+    "backup restore strategies",
+    "log rotation configuration",
+    "performance profiling",
+    # General Knowledge
+    "climate change effects",
+    "renewable energy sources",
+    "electric vehicles",
+    "artificial intelligence ethics",
+    "blockchain technology",
+    "quantum computing basics",
+    "space exploration mars",
+    "gene editing crispr",
+    "vaccine development",
+    "economic indicators gdp",
+    "stock market investing",
+    "cryptocurrency trading",
+    "mental health awareness",
+    "nutrition diet tips",
+    "exercise fitness routine",
+    "meditation mindfulness",
+    "sleep hygiene habits",
+    "stress management",
+    "time management productivity",
+    "remote work tips",
+    "team collaboration",
+    "project management agile",
+    "design thinking process",
+    "user experience research",
+    # Short/Ambiguous Queries (important for training)
+    "cache",
+    "proxy",
+    "queue",
+    "mutex",
+    "semaphore",
+    "deadlock",
+    "heap",
+    "stack",
+    "tree",
+    "graph",
+    "hash",
+    "sort",
+    "api",
+    "sdk",
+    "cli",
+    "gui",
+    "orm",
+    "cdn",
+    "auth",
+    "cors",
+    "csrf",
+    "xss",
+    "jwt",
+    "ssh",
+]
+
+PROMPT_TEMPLATE = """Generate search query expansions for: {query}
+
+Output EXACTLY this format (3 lex, 2 vec, 1 hyde):
+lex: keyword phrase 1
+lex: keyword phrase 2
+lex: keyword phrase 3
+vec: natural language search query
+vec: alternative semantic query
+hyde: A specific 2-sentence document passage answering this query.
+
+Output:"""
+
+
+def generate_with_ollama(
+    query: str, model: str = "gemma3:4b", base_url: str = "http://localhost:11434"
+) -> str | None:
+    """Generate query expansion using Ollama API."""
+    try:
+        response = requests.post(
+            f"{base_url}/api/generate",
+            json={
+                "model": model,
+                "prompt": PROMPT_TEMPLATE.format(query=query),
+                "stream": False,
+                "options": {
+                    "temperature": 0.7,
+                    "top_p": 0.9,
+                    "num_predict": 800,  # More tokens for thinking models
+                },
+            },
+            timeout=120,
+        )
+        response.raise_for_status()
+        return response.json().get("response", "").strip()
+    except Exception as e:
+        print(f"Error generating for '{query}': {e}", file=sys.stderr)
+        return None
+
+
+def parse_expansion(output: str) -> list[list[str]] | None:
+    """Parse the model output into structured format."""
+    items = normalize_output_items(parse_output_text(output))
+    lex_count = sum(1 for kind, _ in items if kind == "lex")
+    vec_count = sum(1 for kind, _ in items if kind == "vec")
+    hyde_count = sum(1 for kind, _ in items if kind == "hyde")
+    if lex_count >= 2 and vec_count >= 1 and hyde_count >= 1:
+        return items
+    return None
+
+
+def generate_query_variations(seed: str) -> list[str]:
+    """Generate variations of a seed query."""
+    variations = [seed]
+
+    # Add question forms
+    if not seed.startswith(("how", "what", "why", "when", "where")):
+        variations.append(f"how to {seed}")
+        variations.append(f"what is {seed}")
+
+    # Add context
+    variations.append(f"{seed} tutorial")
+    variations.append(f"{seed} best practices")
+    variations.append(f"{seed} examples")
+
+    return variations
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate training data using Ollama")
+    parser.add_argument(
+        "--output", "-o", default="data/qmd_expansion_ollama.jsonl", help="Output file"
+    )
+    parser.add_argument(
+        "--count", "-n", type=int, default=1000, help="Number of examples to generate"
+    )
+    parser.add_argument("--model", "-m", default="gemma3:4b", help="Ollama model name")
+    parser.add_argument(
+        "--base-url", default="http://localhost:11434", help="Ollama base URL"
+    )
+    parser.add_argument(
+        "--resume", action="store_true", help="Resume from existing file"
+    )
+    args = parser.parse_args()
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Load existing if resuming
+    existing_queries = set()
+    if args.resume and output_path.exists():
+        with open(output_path) as f:
+            for line in f:
+                obj = json.loads(line)
+                existing_queries.add(obj.get("query", obj.get("input", "")).lower())
+        print(
+            f"Resuming with {len(existing_queries)} existing examples", file=sys.stderr
+        )
+
+    # Generate query pool
+    all_queries = []
+    for seed in QUERY_SEEDS:
+        all_queries.extend(generate_query_variations(seed))
+
+    # Shuffle and filter
+    random.shuffle(all_queries)
+    queries_to_process = [q for q in all_queries if q.lower() not in existing_queries]
+
+    print(
+        f"Processing {min(args.count, len(queries_to_process))} queries with {args.model}...",
+        file=sys.stderr,
+    )
+
+    generated = 0
+    errors = 0
+
+    mode = "a" if args.resume else "w"
+    with open(output_path, mode) as f:
+        for i, query in enumerate(queries_to_process):
+            if generated >= args.count:
+                break
+
+            output = generate_with_ollama(query, args.model, args.base_url)
+            if output:
+                parsed = parse_expansion(output)
+                if parsed:
+                    example = {"query": query, "output": parsed}
+                    f.write(json.dumps(example) + "\n")
+                    f.flush()
+                    generated += 1
+
+                    if generated % 10 == 0:
+                        print(
+                            f"Generated {generated}/{args.count} ({errors} errors)",
+                            file=sys.stderr,
+                        )
+                else:
+                    errors += 1
+            else:
+                errors += 1
+
+            # Small delay to avoid overwhelming the API
+            time.sleep(0.1)
+
+    print(f"\nDone! Generated {generated} examples, {errors} errors", file=sys.stderr)
+    print(f"Output: {output_path}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

package/finetune/dataset/prepare_data.py

@@ -0,0 +1,197 @@
+#!/usr/bin/env python3
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "transformers>=4.45.0",
+#     "jinja2",
+# ]
+# ///
+"""Prepare QMD query expansion data for training.
+
+See PROMPT_FORMAT.md for format specification.
+"""
+
+import argparse
+import json
+import random
+import sys
+import os
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from dataset.schema import (
+    normalize_output_items,
+    output_items_to_text,
+    parse_output_text,
+    has_type,
+)
+
+from transformers import AutoTokenizer
+
+_tokenizer = None
+_tokenizer_model = None
+
+
+def get_tokenizer():
+    global _tokenizer, _tokenizer_model
+    model_name = os.environ.get("QMD_BASE_MODEL", "Qwen/Qwen3-1.7B")
+    if _tokenizer is None or _tokenizer_model != model_name:
+        _tokenizer = AutoTokenizer.from_pretrained(model_name)
+        _tokenizer_model = model_name
+    return _tokenizer
+
+
+
+def format_for_training(query_text: str, output_items: list[list[str]]) -> dict:
+    """Format a single example for SFT training using Qwen chat format."""
+    tokenizer = get_tokenizer()
+    output_text = output_items_to_text(output_items)
+
+    # Use /no_think to disable thinking mode - we want direct output
+    messages = [
+        {
+            "role": "user",
+            "content": f"/no_think Expand this search query: {query_text}",
+        },
+        {"role": "assistant", "content": output_text},
+    ]
+
+    # Use tokenizer to generate proper chat format with special tokens
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+
+    # Strip empty <think> tags - we don't want thinking mode
+    # The template adds "<think>\n\n</think>\n\n" which we remove
+    text = text.replace("<think>\n\n</think>\n\n", "")
+
+    return {
+        "text": text,
+        "messages": messages,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare data for training")
+    parser.add_argument(
+        "--input",
+        type=str,
+        default="data/*.jsonl",
+        help="Input JSONL file(s) - supports glob patterns",
+    )
+    parser.add_argument(
+        "--output", type=str, default="data/train", help="Output directory"
+    )
+    parser.add_argument(
+        "--split", type=float, default=0.1, help="Validation split ratio"
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Shuffle seed (default: 42)",
+    )
+    args = parser.parse_args()
+
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Support glob patterns for input
+    import glob
+
+    if "*" in args.input:
+        input_files = sorted(glob.glob(args.input))
+        if not input_files:
+            print(f"Error: No files found matching: {args.input}")
+            exit(1)
+        print(
+            f"Found {len(input_files)} input files: {[Path(f).name for f in input_files]}"
+        )
+    else:
+        input_path = Path(args.input)
+        if not input_path.exists():
+            print(f"Error: Input file not found: {input_path}")
+            exit(1)
+        input_files = [str(input_path)]
+
+    # Load all examples from all input files
+    examples = []
+
+    for input_file in input_files:
+        file_count = 0
+        with open(input_file) as f:
+            for line in f:
+                if line.strip():
+                    ex = json.loads(line)
+
+                    # Normalize legacy format
+                    if "query" not in ex and "input" in ex:
+                        ex["query"] = ex.pop("input")
+                    if isinstance(ex.get("output"), str):
+                        ex["output"] = parse_output_text(ex["output"])
+                    ex["output"] = normalize_output_items(ex.get("output", []))
+
+                    examples.append(ex)
+                    file_count += 1
+        print(f"  {Path(input_file).name}: {file_count} examples")
+
+    print(f"Loaded {len(examples)} examples total")
+
+    # Combine and shuffle
+    all_examples = examples
+    random.seed(args.seed)
+    random.shuffle(all_examples)
+
+    # Format for training
+    formatted = [format_for_training(ex["query"], ex["output"]) for ex in all_examples]
+
+    # Split into train/val
+    split_idx = int(len(formatted) * (1 - args.split))
+    train_data = formatted[:split_idx]
+    val_data = formatted[split_idx:]
+
+    # Write train set
+    train_path = output_dir / "train.jsonl"
+    with open(train_path, "w") as f:
+        for item in train_data:
+            f.write(json.dumps(item) + "\n")
+
+    # Write validation set
+    val_path = output_dir / "val.jsonl"
+    with open(val_path, "w") as f:
+        for item in val_data:
+            f.write(json.dumps(item) + "\n")
+
+    # Write chat format (for TRL)
+    chat_path = output_dir / "train_chat.jsonl"
+    with open(chat_path, "w") as f:
+        for item in train_data:
+            f.write(json.dumps({"messages": item["messages"]}) + "\n")
+
+    # Stats
+    short_final = sum(1 for ex in all_examples if len(ex["query"].split()) <= 2)
+
+    print(f"\n=== Summary ===")
+    print(f"Total examples: {len(all_examples)}")
+    print(
+        f"Short queries: {short_final} ({100 * short_final / len(all_examples):.1f}%)"
+    )
+    print(f"Train: {len(train_data)}, Val: {len(val_data)}")
+    print(f"Output: {output_dir}")
+
+    # Dataset info
+    dataset_info = {
+        "dataset_name": "qmd-query-expansion",
+        "train_samples": len(train_data),
+        "val_samples": len(val_data),
+        "short_query_pct": round(100 * short_final / len(all_examples), 1),
+        "columns": ["prompt", "completion", "text", "messages"],
+    }
+    with open(output_dir / "dataset_info.json", "w") as f:
+        json.dump(dataset_info, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

package/finetune/dataset/schema.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""Schema helpers for QMD training JSONL data."""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+VALID_OUTPUT_TYPES = {"hyde", "lex", "vec"}
+
+
+def parse_output_text(text: str) -> list[list[str]]:
+    """Parse prefixed output text into list pairs.
+
+    Returns: [["hyde", "..."], ["lex", "..."], ...]
+    """
+    items: list[list[str]] = []
+    for raw_line in text.strip().split("\n"):
+        line = raw_line.strip()
+        if not line:
+            continue
+        if line.startswith("lex:"):
+            items.append(["lex", line[4:].strip()])
+        elif line.startswith("vec:"):
+            items.append(["vec", line[4:].strip()])
+        elif line.startswith("hyde:"):
+            items.append(["hyde", line[5:].strip()])
+    return items
+
+
+def output_items_to_text(items: Iterable[Iterable[str]]) -> str:
+    """Render output list pairs to prefixed text lines."""
+    lines = []
+    for item in items:
+        if not item:
+            continue
+        try:
+            kind, text = item[0], item[1]
+        except Exception:
+            continue
+        if kind not in VALID_OUTPUT_TYPES:
+            continue
+        if text is None:
+            continue
+        text = str(text).strip()
+        if not text:
+            continue
+        lines.append(f"{kind}: {text}")
+    return "\n".join(lines)
+
+
+def normalize_output_items(items: Iterable[Iterable[str]]) -> list[list[str]]:
+    """Normalize output list pairs (filter invalid, trim whitespace)."""
+    normalized: list[list[str]] = []
+    for item in items:
+        if not item:
+            continue
+        try:
+            kind, text = item[0], item[1]
+        except Exception:
+            continue
+        if kind not in VALID_OUTPUT_TYPES:
+            continue
+        if text is None:
+            continue
+        text = str(text).strip()
+        if not text:
+            continue
+        normalized.append([kind, text])
+    return normalized
+
+
+def has_type(items: Iterable[Iterable[str]], kind: str) -> bool:
+    return any(item and item[0] == kind for item in items)

package/finetune/dataset/score_data.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""Score JSONL datasets with the reward function."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+from dataset.schema import (
+    normalize_output_items,
+    output_items_to_text,
+    parse_output_text,
+)
+from reward import score_expansion_detailed
+
+
+def score_file(path: Path) -> tuple[int, int, list[float], dict]:
+    total = 0
+    errors = 0
+    scores: list[float] = []
+    ratings: dict[str, int] = {}
+
+    with path.open("r", encoding="utf-8") as f:
+        for line_num, line in enumerate(f, 1):
+            line = line.strip()
+            if not line:
+                continue
+            total += 1
+            try:
+                obj = json.loads(line)
+            except json.JSONDecodeError:
+                errors += 1
+                continue
+
+            query = obj.get("query") or obj.get("input")
+            output = obj.get("output")
+            if not isinstance(query, str) or not query.strip():
+                errors += 1
+                continue
+            if output is None:
+                errors += 1
+                continue
+
+            if isinstance(output, str):
+                output_items = normalize_output_items(parse_output_text(output))
+            else:
+                output_items = normalize_output_items(output)
+
+            output_text = output_items_to_text(output_items)
+            if not output_text:
+                errors += 1
+                continue
+
+            detail = score_expansion_detailed(query, output_text)
+            score = detail["percentage"]
+            scores.append(score)
+            rating = detail["rating"]
+            ratings[rating] = ratings.get(rating, 0) + 1
+
+    return total, errors, scores, ratings
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Score QMD datasets")
+    parser.add_argument(
+        "paths",
+        nargs="*",
+        default=["finetune/data/*.jsonl"],
+        help="JSONL files or glob patterns (default: finetune/data/*.jsonl)",
+    )
+    args = parser.parse_args()
+
+    repo_root = Path(__file__).parent.parent.parent
+    files: list[Path] = []
+    for pattern in args.paths:
+        if "*" in pattern:
+            files.extend(repo_root.glob(pattern))
+        else:
+            files.append(repo_root / pattern)
+
+    files = [p for p in files if p.exists()]
+    if not files:
+        print("No files found to score.")
+        return 1
+
+    for path in sorted(files):
+        total, errors, scores, ratings = score_file(path)
+        if scores:
+            avg = statistics.mean(scores)
+            median = statistics.median(scores)
+            min_score = min(scores)
+            max_score = max(scores)
+            above_70 = sum(1 for s in scores if s >= 70.0)
+            pct_70 = above_70 / len(scores) * 100
+            print(
+                f"{path}: {len(scores)} scored, {errors} errors, "
+                f"avg {avg:.1f}, median {median:.1f}, min {min_score:.1f}, "
+                f"max {max_score:.1f}, >=70 {pct_70:.1f}%"
+            )
+        else:
+            print(f"{path}: 0 scored, {errors} errors")
+
+        if ratings:
+            rating_parts = [f"{k}:{v}" for k, v in sorted(ratings.items())]
+            print(f"  ratings: {', '.join(rating_parts)}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())