microsoft/TypeAgent
Publicmirrored from https://github.com/microsoft/TypeAgentAvailable
python/fineTuning/unsloth/baseExtract.py
165lines · modecode
| 1 | from argparse import ArgumentParser |
| 2 | import sys |
| 3 | parser = ArgumentParser(description="Extract keywords from dataset using NLTK-RAKE.") |
| 4 | parser.add_argument("--dataset_path", type=str, default='/data/npr/npr_chunks_no_embedding.json', |
| 5 | help="Path to the dataset file.") |
| 6 | # output_file |
| 7 | parser.add_argument("--output_file", type=str, default='baseExtraction.txt', |
| 8 | help="Path to the output file.") |
| 9 | args = parser.parse_args(sys.argv[1:]) |
| 10 | dataset_path = args.dataset_path |
| 11 | output_file = args.output_file |
| 12 | |
| 13 | import json |
| 14 | import re |
| 15 | import time |
| 16 | from dataclasses import dataclass, field |
| 17 | import matplotlib.pyplot as plt |
| 18 | |
| 19 | @dataclass |
| 20 | class Message: |
| 21 | content: str |
| 22 | speaker: str |
| 23 | section: 'Section' = None |
| 24 | |
| 25 | @dataclass |
| 26 | class Section: |
| 27 | title: str |
| 28 | messages: list[Message] = field(default_factory=list) |
| 29 | words: set[str] = field(default_factory=set) |
| 30 | total_words: int = 0 |
| 31 | |
| 32 | @dataclass |
| 33 | class Word: |
| 34 | word: str |
| 35 | messages: list[Message] = field(default_factory=list) |
| 36 | |
| 37 | def read_and_count_messages(file_path): |
| 38 | """Read the JSON file and return the array and its count.""" |
| 39 | with open(file_path, 'r') as f: |
| 40 | data = json.load(f) |
| 41 | |
| 42 | count = len(data) |
| 43 | print(f"Total messages in dataset: {count}") |
| 44 | |
| 45 | return data, count |
| 46 | |
| 47 | # Start timing |
| 48 | start_time = time.time() |
| 49 | |
| 50 | # Read the dataset |
| 51 | messages, message_count = read_and_count_messages(dataset_path) |
| 52 | |
| 53 | # Create dictionary of sections and process messages |
| 54 | print("Processing messages and extracting words...") |
| 55 | word_pattern = re.compile(r'\b\w+\b') |
| 56 | sections = {} |
| 57 | all_messages = [] |
| 58 | word_dict = {} |
| 59 | total_words = 0 |
| 60 | all_unique_words = set() |
| 61 | |
| 62 | for i, message in enumerate(messages): |
| 63 | section_title = message['section_title'] |
| 64 | if section_title not in sections: |
| 65 | sections[section_title] = Section(title=section_title) |
| 66 | |
| 67 | content = message['content'] |
| 68 | speaker = message['speaker'] |
| 69 | msg = Message(content=content, speaker=speaker, section=sections[section_title]) |
| 70 | sections[section_title].messages.append(msg) |
| 71 | all_messages.append(msg) |
| 72 | |
| 73 | # Extract words from this message |
| 74 | words = word_pattern.findall(content.lower()) |
| 75 | word_count = len(words) |
| 76 | total_words += word_count |
| 77 | sections[section_title].total_words += word_count |
| 78 | sections[section_title].words.update(words) |
| 79 | all_unique_words.update(words) |
| 80 | |
| 81 | # Track which messages contain each word |
| 82 | unique_message_words = set(words) |
| 83 | for word in unique_message_words: |
| 84 | if word not in word_dict: |
| 85 | word_dict[word] = Word(word=word) |
| 86 | word_dict[word].messages.append(msg) |
| 87 | |
| 88 | print(f"Processing complete.") |
| 89 | |
| 90 | # Print statistics |
| 91 | section_count = len(sections) |
| 92 | total_messages = sum(len(section.messages) for section in sections.values()) |
| 93 | avg_messages_per_section = total_messages / section_count if section_count > 0 else 0 |
| 94 | avg_words_per_section = sum(len(section.words) for section in sections.values()) / section_count if section_count > 0 else 0 |
| 95 | |
| 96 | # Calculate word density ratio for each section |
| 97 | word_density_ratios = [] |
| 98 | for section in sections.values(): |
| 99 | if section.total_words > 0: |
| 100 | ratio = len(section.words) / section.total_words |
| 101 | word_density_ratios.append(ratio) |
| 102 | |
| 103 | avg_word_density = sum(word_density_ratios) / len(word_density_ratios) if word_density_ratios else 0 |
| 104 | |
| 105 | total_message_references = sum(len(word.messages) for word in word_dict.values()) |
| 106 | avg_messages_per_word = total_message_references / len(word_dict) if word_dict else 0 |
| 107 | |
| 108 | print(f"Total sections: {section_count}") |
| 109 | print(f"Average messages per section: {avg_messages_per_section:.2f}") |
| 110 | print(f"Total words: {total_words}") |
| 111 | print(f"Unique words across all messages: {len(all_unique_words)}") |
| 112 | print(f"Average unique words per section: {avg_words_per_section:.2f}") |
| 113 | print(f"Average word density ratio (unique/total per section): {avg_word_density * 100:.3f}%") |
| 114 | print(f"Overall word density ratio (unique/total across all messages): {len(all_unique_words) / total_words * 100:.3f}%") |
| 115 | print(f"Average messages per unique word: {avg_messages_per_word:.2f}") |
| 116 | |
| 117 | # Analyze message count distribution |
| 118 | message_counts = [len(word.messages) for word in word_dict.values()] |
| 119 | message_counts_sorted = sorted(message_counts, reverse=True) |
| 120 | |
| 121 | # Get top 5 words by message count |
| 122 | top_words = sorted(word_dict.values(), key=lambda w: len(w.messages), reverse=True)[:5] |
| 123 | print(f"\nTop 5 words by message count:") |
| 124 | for i, word_obj in enumerate(top_words, 1): |
| 125 | print(f" {i}. '{word_obj.word}': {len(word_obj.messages)} messages") |
| 126 | |
| 127 | # Calculate statistics |
| 128 | import statistics |
| 129 | median_messages = statistics.median(message_counts) if message_counts else 0 |
| 130 | max_messages = max(message_counts) if message_counts else 0 |
| 131 | min_messages = min(message_counts) if message_counts else 0 |
| 132 | |
| 133 | print(f"\nMessage count statistics per word:") |
| 134 | print(f" Min: {min_messages}") |
| 135 | print(f" Median: {median_messages:.2f}") |
| 136 | print(f" Max: {max_messages}") |
| 137 | |
| 138 | # Count words appearing in 4 or fewer messages |
| 139 | words_4_or_fewer = sum(1 for count in message_counts if count <= 4) |
| 140 | percentage_4_or_fewer = (words_4_or_fewer / len(message_counts) * 100) if message_counts else 0 |
| 141 | print(f" Words appearing in ≤4 messages: {words_4_or_fewer} ({percentage_4_or_fewer:.2f}%)") |
| 142 | |
| 143 | # Count words appearing in 16 or fewer messages |
| 144 | words_16_or_fewer = sum(1 for count in message_counts if count <= 16) |
| 145 | percentage_16_or_fewer = (words_16_or_fewer / len(message_counts) * 100) if message_counts else 0 |
| 146 | print(f" Words appearing in ≤16 messages: {words_16_or_fewer} ({percentage_16_or_fewer:.2f}%)") |
| 147 | |
| 148 | # Plot distribution of messages per word |
| 149 | print("\nCreating distribution plot...") |
| 150 | message_counts = [len(word.messages) for word in word_dict.values()] |
| 151 | |
| 152 | plt.figure(figsize=(12, 6)) |
| 153 | plt.hist(message_counts, bins=50, edgecolor='black', alpha=0.7) |
| 154 | plt.xlabel('Number of Messages per Word') |
| 155 | plt.ylabel('Number of Words') |
| 156 | plt.title('Distribution of Message Counts per Unique Word') |
| 157 | plt.yscale('log') # Use log scale for y-axis since distribution is likely skewed |
| 158 | plt.grid(True, alpha=0.3) |
| 159 | plt.tight_layout() |
| 160 | plt.savefig('word_message_distribution.png', dpi=150) |
| 161 | print(f"Distribution plot saved to word_message_distribution.png") |
| 162 | |
| 163 | # Print elapsed time |
| 164 | elapsed_time = time.time() - start_time |
| 165 | print(f"\nTotal processing time: {elapsed_time:.2f} seconds") |