import json import os from pathlib import Path import pandas as pd from typing import List, Dict, Tuple, Optional import random from tqdm import tqdm import re import matplotlib.pyplot as plt import numpy as np class KokoroChatPreprocessor: def __init__(self, data_path: str, max_length: int = 2048, min_score: int = 60): """ Initialize the preprocessor for KokoroChat dataset Args: data_path: Path to KokoroChat repository max_length: Maximum sequence length for model input min_score: Minimum score threshold for filtering conversations (default: 60) """ self.data_path = Path(data_path) self.max_length = max_length self.min_score = min_score self.conversations = [] self.score_distribution = [] # Track score distribution self.system_prompt = """あなたは思いやりのある心理カウンセラーです。 クライアントの感情を理解し、共感的で支援的な応答を提供してください。 プライバシーを尊重し、判断を下さず、希望と実用的な洞察を提供することに焦点を当ててください。""" def load_json_files(self) -> List[Dict]: """Load all JSON files from the dataset""" json_files = [] # Changed from "data" to "kokorochat_dialogues" data_dir = self.data_path / "kokorochat_dialogues" # Check if data directory exists, if not try root directory if not data_dir.exists(): data_dir = self.data_path print(f"Using root directory: {data_dir}") else: print(f"Using data directory: {data_dir}") for root, dirs, files in os.walk(data_dir): for file in tqdm(files, desc="Loading JSON files"): if file.endswith('.json'): file_path = os.path.join(root, file) try: with open(file_path, 'r', encoding='utf-8') as f: data = json.load(f) json_files.append(data) except Exception as e: print(f"Error loading {file_path}: {e}") return json_files def analyze_score_distribution(self, json_files: List[Dict]) -> Dict: """ Analyze the distribution of scores in the dataset Returns: Dictionary with score statistics """ scores = [] for data in json_files: if 'review_by_client_jp' in data: score = data['review_by_client_jp'].get('点数', 0) if score > 0: # Only count valid scores scores.append(score) self.score_distribution.append(score) if scores: stats = { 'total_conversations': len(json_files), 'conversations_with_scores': len(scores), 'mean_score': float(np.mean(scores)), 'median_score': float(np.median(scores)), 'std_score': float(np.std(scores)), 'min_score': float(np.min(scores)), 'max_score': float(np.max(scores)), 'percentiles': { '25th': float(np.percentile(scores, 25)), '50th': float(np.percentile(scores, 50)), '75th': float(np.percentile(scores, 75)), '90th': float(np.percentile(scores, 90)) }, 'score_ranges': { '0-30': int(sum(1 for s in scores if 0 <= s < 30)), '30-50': int(sum(1 for s in scores if 30 <= s < 50)), '50-60': int(sum(1 for s in scores if 50 <= s < 60)), '60-70': int(sum(1 for s in scores if 60 <= s < 70)), '70-80': int(sum(1 for s in scores if 70 <= s < 80)), '80-90': int(sum(1 for s in scores if 80 <= s < 90)), '90-100': int(sum(1 for s in scores if 90 <= s <= 100)), } } # Calculate how many conversations would be kept at different thresholds threshold_analysis = {} for threshold in [30, 40, 50, 60, 65, 70, 75, 80]: kept = sum(1 for s in scores if s >= threshold) threshold_analysis[f'threshold_{threshold}'] = { 'conversations_kept': kept, 'percentage_kept': round((kept / len(scores)) * 100, 2) } stats['threshold_analysis'] = threshold_analysis return stats else: return {'error': 'No valid scores found in dataset'} def plot_score_distribution(self, save_path: str = "score_distribution.png"): """ Plot the distribution of scores """ if not self.score_distribution: print("No scores to plot. Run analyze_score_distribution first.") return fig, axes = plt.subplots(2, 2, figsize=(15, 10)) # Histogram axes[0, 0].hist(self.score_distribution, bins=20, edgecolor='black', alpha=0.7) axes[0, 0].axvline(self.min_score, color='red', linestyle='--', label=f'Current threshold: {self.min_score}') axes[0, 0].set_xlabel('Score') axes[0, 0].set_ylabel('Frequency') axes[0, 0].set_title('Score Distribution') axes[0, 0].legend() axes[0, 0].grid(True, alpha=0.3) # Box plot axes[0, 1].boxplot(self.score_distribution, vert=True) axes[0, 1].set_ylabel('Score') axes[0, 1].set_title('Score Box Plot') axes[0, 1].grid(True, alpha=0.3) # Cumulative distribution sorted_scores = np.sort(self.score_distribution) cumulative = np.arange(1, len(sorted_scores) + 1) / len(sorted_scores) axes[1, 0].plot(sorted_scores, cumulative) axes[1, 0].axvline(self.min_score, color='red', linestyle='--', label=f'Current threshold: {self.min_score}') axes[1, 0].set_xlabel('Score') axes[1, 0].set_ylabel('Cumulative Probability') axes[1, 0].set_title('Cumulative Distribution') axes[1, 0].legend() axes[1, 0].grid(True, alpha=0.3) # Threshold impact analysis thresholds = range(30, 90, 5) kept_percentages = [] for t in thresholds: kept = sum(1 for s in self.score_distribution if s >= t) kept_percentages.append((kept / len(self.score_distribution)) * 100) axes[1, 1].plot(thresholds, kept_percentages, marker='o') axes[1, 1].axvline(self.min_score, color='red', linestyle='--', label=f'Current threshold: {self.min_score}') axes[1, 1].set_xlabel('Score Threshold') axes[1, 1].set_ylabel('% of Conversations Kept') axes[1, 1].set_title('Impact of Score Threshold') axes[1, 1].legend() axes[1, 1].grid(True, alpha=0.3) plt.tight_layout() plt.savefig(save_path, dpi=300, bbox_inches='tight') plt.show() print(f"Score distribution plot saved to {save_path}") def extract_high_quality_conversations(self, data: Dict) -> List[Dict]: """ Extract conversations with high counselor ratings based on min_score Focus on conversations where counselor performed well """ conversations = [] # Check if review exists and has good score if 'review_by_client_jp' in data: review = data['review_by_client_jp'] score = review.get('点数', 0) # Use configurable min_score threshold if score >= self.min_score: dialogue = data.get('dialogue', []) # Create conversation pairs conversation_text = "" for turn in dialogue: role = turn['role'] utterance = turn['utterance'] if role == 'counselor': conversation_text += f"カウンセラー: {utterance}\n" else: conversation_text += f"クライアント: {utterance}\n" # Extract detailed metrics for potential weighted training conversations.append({ 'text': conversation_text, 'score': score, # Store the score here 'topic': data.get('topic', {}).get('main_jp', 'Unknown'), 'review_metrics': { 'empathy': review.get('聴いてもらえた、わかってもらえたと感じた', 0), 'respect': review.get('尊重されたと感じた', 0), 'insights': review.get('新しい気づきや体験があった', 0), 'hope': review.get('希望や期待を感じられた', 0), 'concerns_addressed': review.get('取り組みたかったことを扱えた', 0), 'collaboration': review.get('一緒に考えながら取り組めた', 0), 'rhythm': review.get('やりとりのリズムがあっていた', 0), 'comfort': review.get('居心地のよいやりとりだった', 0), 'overall_appropriate': review.get('全体として適切でよかった', 0), 'valuable': review.get('今回の相談は価値があった', 0), 'smooth_start': review.get('相談開始の円滑さ', 0), 'good_ending': review.get('相談終了のタイミング（不必要に聴きすぎていないか）、円滑さ', 0), 'acceptance_empathy': review.get('受容·共感', 0), 'affirmation': review.get('肯定·承認', 0), 'effective_questions': review.get('的確な質問による会話の促進', 0), 'summarization': review.get('要約', 0), 'problem_clarification': review.get('問題の明確化', 0), 'goal_clarification': review.get('この相談での目標の明確化', 0), 'actionable_suggestions': review.get('次の行動につながる提案', 0), 'encouragement': review.get('勇気づけ·希望の喚起', 0) } }) return conversations def create_training_examples(self, conversations: List[Dict], use_weighted_sampling: bool = False) -> List[Dict]: """ Create training examples in instruction-following format Args: conversations: List of conversation dictionaries use_weighted_sampling: If True, create more examples from higher-scored conversations """ training_examples = [] for conv in tqdm(conversations, desc="Creating training examples"): dialogue_lines = conv['text'].split('\n') score = conv['score'] # Get score from the conversation dict # Calculate sampling weight based on score if enabled if use_weighted_sampling: # Higher scores get more weight (normalized to 1-3 range) weight = max(1, int((score - self.min_score) / 20) + 1) else: weight = 1 # Create multiple training examples from each conversation for _ in range(weight): # Repeat based on weight for i in range(0, len(dialogue_lines) - 1, 2): if i + 1 < len(dialogue_lines): client_line = dialogue_lines[i] counselor_line = dialogue_lines[i + 1] # Check if lines contain the expected prefixes if 'クライアント:' in client_line and 'カウンセラー:' in counselor_line: client_msg = client_line.replace('クライアント: ', '').replace('クライアント:', '').strip() counselor_msg = counselor_line.replace('カウンセラー: ', '').replace('カウンセラー:', '').strip() # Skip empty messages if not client_msg or not counselor_msg: continue # Format for instruction tuning example = { 'instruction': self.system_prompt, 'input': client_msg, 'output': counselor_msg, 'score': score, # Use the score from conversation 'topic': conv['topic'], 'metrics': conv['review_metrics'] # Include detailed metrics } training_examples.append(example) return training_examples def prepare_dataset(self, test_size: float = 0.1, val_size: float = 0.1, use_weighted_sampling: bool = False, analyze_scores: bool = True): """ Prepare train, validation, and test datasets Args: test_size: Proportion of data for testing val_size: Proportion of data for validation use_weighted_sampling: If True, oversample high-quality conversations analyze_scores: If True, print score distribution analysis """ print("Loading KokoroChat dataset...") json_files = self.load_json_files() print(f"Loaded {len(json_files)} conversation files") # Analyze score distribution if requested if analyze_scores: print("\n" + "="*60) print("SCORE DISTRIBUTION ANALYSIS") print("="*60) stats = self.analyze_score_distribution(json_files) if 'error' not in stats: print(f"Total conversations: {stats['total_conversations']}") print(f"Conversations with scores: {stats['conversations_with_scores']}") print(f"\nScore Statistics:") print(f" Mean: {stats['mean_score']:.2f}") print(f" Median: {stats['median_score']:.2f}") print(f" Std Dev: {stats['std_score']:.2f}") print(f" Range: {stats['min_score']:.0f} - {stats['max_score']:.0f}") print(f"\nScore Distribution:") for range_name, count in stats['score_ranges'].items(): percentage = (count / stats['conversations_with_scores']) * 100 print(f" {range_name}: {count} ({percentage:.1f}%)") print(f"\nThreshold Impact Analysis:") for threshold_name, data in stats['threshold_analysis'].items(): threshold = threshold_name.split('_')[1] print(f" Threshold >= {threshold}: {data['conversations_kept']} conversations ({data['percentage_kept']:.1f}%)") print(f"\nCurrent threshold ({self.min_score}) will keep: ", end="") kept = sum(1 for s in self.score_distribution if s >= self.min_score) print(f"{kept} conversations ({(kept/len(self.score_distribution))*100:.1f}%)") print("="*60 + "\n") # Plot distribution self.plot_score_distribution() all_conversations = [] filtered_count = 0 total_count = 0 for data in json_files: if 'review_by_client_jp' in data: total_count += 1 score = data['review_by_client_jp'].get('点数', 0) if score < self.min_score: filtered_count += 1 conversations = self.extract_high_quality_conversations(data) all_conversations.extend(conversations) print(f"Filtered out {filtered_count} conversations with score < {self.min_score}") print(f"Extracted {len(all_conversations)} high-quality conversations (score >= {self.min_score})") # Create training examples training_examples = self.create_training_examples( all_conversations, use_weighted_sampling=use_weighted_sampling ) print(f"Created {len(training_examples)} training examples") if use_weighted_sampling: print("Note: Used weighted sampling - higher scored conversations appear more frequently") # Shuffle and split random.shuffle(training_examples) total_size = len(training_examples) test_split = int(total_size * test_size) val_split = int(total_size * val_size) test_data = training_examples[:test_split] val_data = training_examples[test_split:test_split + val_split] train_data = training_examples[test_split + val_split:] print(f"\nDataset splits:") print(f" Train: {len(train_data)} examples") print(f" Validation: {len(val_data)} examples") print(f" Test: {len(test_data)} examples") return { 'train': train_data, 'validation': val_data, 'test': test_data } def format_for_lfm(self, example: Dict) -> str: """ Format example for LFM model training """ formatted = f"""### Instruction: {example['instruction']} ### Input: {example['input']} ### Response: {example['output']}""" return formatted def save_datasets(self, datasets: Dict, output_dir: str): """Save processed datasets with proper type conversion for JSON serialization""" output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) # Helper function to convert numpy types to Python native types def convert_to_native(obj): if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() else: return obj # Save dataset statistics stats = { 'min_score_threshold': int(self.min_score), 'dataset_sizes': { 'train': len(datasets['train']), 'validation': len(datasets['validation']), 'test': len(datasets['test']) }, 'score_distribution': {} } for split_name, data in datasets.items(): # Calculate score distribution for this split scores = [ex['score'] for ex in data] if scores: stats['score_distribution'][split_name] = { 'mean': float(np.mean(scores)), 'median': float(np.median(scores)), 'min': float(np.min(scores)), 'max': float(np.max(scores)), 'std': float(np.std(scores)) } # Save as JSONL for easier streaming file_path = output_path / f"{split_name}.jsonl" with open(file_path, 'w', encoding='utf-8') as f: for example in data: formatted_text = self.format_for_lfm(example) # Convert all numpy types to native Python types json_obj = { 'text': formatted_text, 'score': convert_to_native(example['score']), 'topic': example['topic'] } json_line = json.dumps(json_obj, ensure_ascii=False) f.write(json_line + '\n') print(f"Saved {split_name} dataset with {len(data)} examples to {file_path}") # Save statistics stats_path = output_path / "dataset_stats.json" with open(stats_path, 'w', encoding='utf-8') as f: json.dump(stats, f, ensure_ascii=False, indent=2) print(f"Saved dataset statistics to {stats_path}") # Print summary statistics print("\n" + "="*60) print("DATASET SUMMARY") print("="*60) print(f"Minimum score threshold: {stats['min_score_threshold']}") print("\nDataset sizes:") for split, size in stats['dataset_sizes'].items(): print(f" {split}: {size} examples") print("\nScore distributions by split:") for split, dist in stats['score_distribution'].items(): print(f" {split}:") print(f" Mean: {dist['mean']:.2f}") print(f" Std: {dist['std']:.2f}") print(f" Range: {dist['min']:.0f} - {dist['max']:.0f}") print("="*60) # Run preprocessing with different score thresholds if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description='Preprocess KokoroChat dataset') parser.add_argument('--data_path', type=str, default='./KokoroChat', help='Path to KokoroChat repository') parser.add_argument('--min_score', type=int, default=70, help='Minimum score threshold for filtering (default: 70)') parser.add_argument('--output_dir', type=str, default='./processed_data', help='Output directory for processed data') parser.add_argument('--weighted_sampling', action='store_true', help='Use weighted sampling based on scores') parser.add_argument('--test_size', type=float, default=0.1, help='Test set size (default: 0.1)') parser.add_argument('--val_size', type=float, default=0.1, help='Validation set size (default: 0.1)') parser.add_argument('--analyze_only', action='store_true', help='Only analyze score distribution without processing') args = parser.parse_args() # Initialize preprocessor with configurable min_score preprocessor = KokoroChatPreprocessor( data_path=args.data_path, min_score=args.min_score ) if args.analyze_only: # Just analyze the score distribution print("Running score distribution analysis only...") json_files = preprocessor.load_json_files() stats = preprocessor.analyze_score_distribution(json_files) preprocessor.plot_score_distribution(f"score_analysis_threshold_{args.min_score}.png") else: # Full preprocessing print(f"Processing with minimum score threshold: {args.min_score}") datasets = preprocessor.prepare_dataset( test_size=args.test_size, val_size=args.val_size, use_weighted_sampling=args.weighted_sampling, analyze_scores=True ) # Save with threshold in directory name output_dir = f"{args.output_dir}_score{args.min_score}" preprocessor.save_datasets(datasets, output_dir) print(f"\nProcessing complete! Data saved to {output_dir}") print("\nNext steps:") print("1. Run fine-tuning: python finetune_lfm.py") print("2. Run benchmarking: python benchmark_model.py") print("3. Optimize for mobile: python optimize_for_mobile.py")