topic-tree.py

from promptsmith import EngineArguments, DataEngine, TopicTree, TopicTreeArguments

# Define the system prompt
system_prompt = (
    "You are a knowledgeable dog expert. You provide detailed advice about dogs, "
    "including training, health, breeds, and care. Always include specific examples "
    "and practical advice."
)

# Create root prompts based on major dog topics
tree = TopicTree(
    args=TopicTreeArguments(
        root_prompt="Dog Knowledge and Care",
        model_system_prompt=system_prompt,
        tree_degree=5,    # 5 subtopics per branch
        tree_depth=3      # 3 levels deep
    )
)

# This would generate a tree like:
"""
Dog Knowledge and Care
├── Training
│   ├── Basic Commands
│   │   ├── Sit
│   │   ├── Stay
│   │   ├── Come
│   ├── Behavioral Issues
│   │   ├── Barking
│   │   ├── Chewing
│   │   ├── Jumping
├── Health Care
│   ├── Preventive Care
│   │   ├── Vaccinations
│   │   ├── Regular Checkups
│   ├── Common Issues
│   │   ├── Skin Problems
│   │   ├── Digestive Issues
└── Breeds
    ├── Small Dogs
    │   ├── Chihuahua
    │   ├── Yorkshire Terrier
    ├── Large Dogs
    │   ├── German Shepherd
    │   ├── Labrador Retriever
"""

tree.build_tree(model_name="llama3:latest")
tree.save("dog_knowledge_tree.jsonl")

# Initialize the DataEngine
engine = DataEngine(
    args=EngineArguments(
        instructions=(
            "Generate question and answer pairs about dogs. Each response should be detailed "
            "and specific to the topic. Include practical examples and real-world applications."
        ),
        system_prompt=system_prompt,
    )
)

# Generate dataset using the topic tree
dataset = engine.create_data(
    model_name="llama3:latest",
    num_steps=20,
    batch_size=5,
    topic_tree=tree
)

# This would generate samples like:
"""
{
    "messages": [
        {
            "role": "user",
            "content": "What's the best way to train a German Shepherd to stop excessive barking?"
        },
        {
            "role": "assistant",
            "content": "Training a German Shepherd to control barking requires consistency and positive reinforcement. Here's a step-by-step approach:

1. Identify triggers: First, determine what causes the excessive barking (strangers, other dogs, noises).

2. Teach the 'quiet' command:
   - When your dog barks, wait for a brief pause
   - Immediately say 'quiet' in a firm but calm voice
   - When they remain quiet, reward with a treat and praise
   
3. Use redirection:
   - When you anticipate a barking trigger, redirect attention
   - Ask for another behavior like 'sit' or 'look at me'
   - Reward compliance

Remember that German Shepherds are naturally protective, so some barking is normal. The goal is to teach them when it's appropriate.

Tip: Never use punishment as it can increase anxiety and make barking worse."
        }
    ]
}
"""

# Save the dataset
dataset.save("dog_training_dataset.jsonl")