Offline Training (SFT/DPO)

1
def filter_rollouts(input_file: str, output_file: str, filters: Dict):
2
    """Apply automatic quality filters to rollouts."""
3
    with open(input_file) as f, open(output_file, 'w') as out:
4
        kept = 0
5
        total = 0
6
        
7
        for line in f:
8
            rollout = json.loads(line)
9
            total += 1
10
            
11
            # Apply filters
12
            if (rollout.get('reward', 0) >= filters.get('min_reward', 0.5) and
13
                rollout.get('success', False) and
14
                len(rollout.get('output', [])) >= filters.get('min_turns', 2) and
15
                len(rollout.get('output', [])) <= filters.get('max_turns', 20)):
16
                
17
                out.write(line)
18
                kept += 1
19
        
20
        print(f"Kept {kept}/{total} rollouts ({kept/total*100:.1f}%)")
21
22
# Apply filters first
23
filter_rollouts('raw_rollouts.jsonl', 'filtered_rollouts.jsonl', {
24
    'min_reward': 0.7,
25
    'min_turns': 3,
26
    'max_turns': 15
27
})

1
def sample_for_review(input_file: str, sample_size: int = 50):
2
    """Sample rollouts for manual review."""
3
    import random
4
    
5
    with open(input_file) as f:
6
        rollouts = [json.loads(line) for line in f]
7
    
8
    # Stratified sampling by reward
9
    low_reward = [r for r in rollouts if r.get('reward', 0) < 0.5]
10
    mid_reward = [r for r in rollouts if 0.5 <= r.get('reward', 0) < 0.8]
11
    high_reward = [r for r in rollouts if r.get('reward', 0) >= 0.8]
12
    
13
    sample = (random.sample(low_reward, min(10, len(low_reward))) +
14
              random.sample(mid_reward, min(20, len(mid_reward))) +
15
              random.sample(high_reward, min(20, len(high_reward))))
16
    
17
    with open('manual_review_sample.jsonl', 'w') as out:
18
        for rollout in sample:
19
            out.write(json.dumps(rollout) + '\n')

1
import json
2
from typing import List, Dict
3
4
def process_sft_data(filtered_rollout_file: str, output_file: str):
5
    """Convert filtered rollouts to SFT training format."""
6
    with open(filtered_rollout_file) as f, open(output_file, 'w') as out:
7
        for line in f:
8
            rollout = json.loads(line)
9
            sft_example = {
10
                "messages": rollout['output'],
11
                "reward": rollout['reward'],
12
                "task_type": rollout.get('metadata', {}).get('task_type', 'general')
13
            }
14
            out.write(json.dumps(sft_example) + '\n')
15
16
# Process filtered rollouts (no additional filtering needed)
17
process_sft_data('filtered_rollouts.jsonl', 'sft_data.jsonl')

1
def create_dpo_pairs(filtered_rollout_file: str, output_file: str):
2
    """Create preference pairs from pairs of filtered rollouts."""
3
    
4
    # Group rollouts by task
5
    task_groups = {}
6
    with open(filtered_rollout_file) as f:
7
        for line in f:
8
            rollout = json.loads(line)
9
            task_id = rollout.get('task_id') or hash(json.dumps(rollout['responses_create_params']['input']))
10
            
11
            if task_id not in task_groups:
12
                task_groups[task_id] = []
13
            task_groups[task_id].append(rollout)
14
    
15
    # Create preference pairs from pairs of rollouts
16
    pairs = []
17
    for task_rollouts in task_groups.values():
18
        if len(task_rollouts) == 2:  # DPO works with pairs
19
            rollout_1, rollout_2 = task_rollouts
20
            
21
            # Determine which is better based on reward
22
            if rollout_1['reward'] > rollout_2['reward']:
23
                chosen, rejected = rollout_1, rollout_2
24
            else:
25
                chosen, rejected = rollout_2, rollout_1
26
            
27
            # Only create pair if there's meaningful difference
28
            quality_diff = chosen['reward'] - rejected['reward']
29
            if quality_diff >= 0.1:  # Minimum difference threshold
30
                pairs.append({
31
                    "prompt": chosen['responses_create_params']['input'],
32
                    "chosen": chosen['output'],
33
                    "rejected": rejected['output'],
34
                    "quality_difference": quality_diff
35
                })
36
    
37
    # Save preference pairs
38
    with open(output_file, 'w') as out:
39
        for pair in pairs:
40
            out.write(json.dumps(pair) + '\n')
41
    
42
    print(f"Created {len(pairs)} preference pairs")
43
44
# Create DPO pairs from filtered rollouts
45
create_dpo_pairs('filtered_rollouts.jsonl', 'dpo_pairs.jsonl')

1
# Prefer high-quality filtered data over large noisy datasets
2
filter_criteria = {
3
    'min_reward': 0.8,        # High threshold for SFT
4
    'min_success_rate': 0.9,
5
    'require_tool_usage': True  # Domain-specific requirements
6
}

1
# Ensure diverse task representation
2
def balance_dataset(input_file: str, output_file: str, max_per_category: int = 100):
3
    task_counts = {}
4
    balanced_data = []
5
    
6
    with open(input_file) as f:
7
        for line in f:
8
            data = json.loads(line)
9
            task_type = data.get('metadata', {}).get('task_type', 'general')
10
            
11
            if task_counts.get(task_type, 0) < max_per_category:
12
                balanced_data.append(data)
13
                task_counts[task_type] = task_counts.get(task_type, 0) + 1
14
    
15
    with open(output_file, 'w') as out:
16
        for data in balanced_data:
17
            out.write(json.dumps(data) + '\n')