optimized time complexity'

bae_chat_log_analysis/__main__.py

@@ -1,62 +1,84 @@
+import argparse
+import cProfile
+import io
+import pstats
+import time
 from pathlib import Path
 
+import tiktoken
+
 from .abbreviations import generate_emoji_abbreviations
 from .reducer import reduce_tokens
-from .token_counter import count_tokens
 
-SOURCE_FILE = Path("bae_chat_logs/Source/WhatsApp_Bae_Chat.txt")
-REDUCED_FILE = Path("bae_chat_logs/Reduced/WhatsApp_Chat-08-03-2026-reduced")
+_DEFAULT_SOURCE = Path("bae_chat_logs/Source/WhatsApp_Bae_Chat.txt")
+_DEFAULT_OUTPUT = Path("bae_chat_logs/Reduced/WhatsApp_Chat-08-03-2026-reduced")
 
-SPLIT_OPTIONS = {'1': 'year', '2': 'month', '3': 'week', '4': 'day', '5': 'none'}
+def _fmt(seconds):
+    return f"{seconds * 1000:.1f}ms" if seconds < 1 else f"{seconds:.2f}s"
 
 
-def prompt_split_option():
-    print("How would you like to split the chat output?")
-    print("  1) By year")
-    print("  2) By month")
-    print("  3) By week")
-    print("  4) By day")
-    print("  5) No split (single file)")
-    while True:
-        choice = input("Enter choice (1-5): ").strip()
-        if choice in SPLIT_OPTIONS:
-            return SPLIT_OPTIONS[choice]
-        print("Invalid choice. Please enter 1, 2, 3, 4, or 5.")
-
-
-def main():
-    split_by = prompt_split_option()
+def main(source_file=_DEFAULT_SOURCE, output_file=_DEFAULT_OUTPUT, split_by="month"):
     print()
 
+    t_total = time.perf_counter()
+
     # 1. Generate abbreviations
     print("Generating emoji abbreviations...")
+    t0 = time.perf_counter()
     abbreviations = generate_emoji_abbreviations()
-    print(f"  {len(abbreviations)} abbreviations loaded")
+    print(f"  {len(abbreviations)} abbreviations loaded  [{_fmt(time.perf_counter() - t0)}]")
 
-    # 2. Reduce source file
-    print(f"Reducing {SOURCE_FILE} (split by: {split_by})...")
-    output_files = reduce_tokens(SOURCE_FILE, REDUCED_FILE, abbreviations, split_by)
+    # 2. Reduce source file (token counts come back from in-memory content — no re-reads)
+    print(f"Reducing {source_file} (split by: {split_by})...")
+    t0 = time.perf_counter()
+    encoding = tiktoken.get_encoding("cl100k_base")
+    output_files, reduced_tokens_total, source_tokens = reduce_tokens(
+        source_file, output_file, abbreviations, split_by, encoding=encoding
+    )
+    print(f"  Reduction complete  [{_fmt(time.perf_counter() - t0)}]")
     for f in output_files:
         print(f"  Saved to {f}")
 
-    # 3. Count and display tokens
-    print("Counting tokens...")
-    source_tokens = count_tokens(SOURCE_FILE)
-
-    reduced_tokens_total = 0
-    for f in output_files:
-        tokens = count_tokens(f)
-        reduced_tokens_total += tokens
-        if len(output_files) > 1:
-            print(f"  {f.name}: {tokens:,} tokens")
-
     saved = source_tokens - reduced_tokens_total
 
     print()
     print(f"Source tokens:  {source_tokens:,}")
     print(f"Reduced tokens: {reduced_tokens_total:,}")
     print(f"Tokens saved:   {saved:,} ({saved / source_tokens * 100:.1f}%)")
+    print(f"\nTotal time: {_fmt(time.perf_counter() - t_total)}")
 
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("source", nargs="?", type=Path, default=_DEFAULT_SOURCE,
+                        help="Source chat log file (default: %(default)s)")
+    parser.add_argument("output", nargs="?", type=Path, default=_DEFAULT_OUTPUT,
+                        help="Output file base path (default: %(default)s)")
+    parser.add_argument(
+        "--split",
+        choices=["year", "month", "week", "day", "none"],
+        default="none",
+        help="How to split the chat output (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--profile",
+        nargs="?",
+        const=20,
+        type=int,
+        metavar="N",
+        help="Run with cProfile and print the top N functions by cumulative time (default: 20)",
+    )
+    args = parser.parse_args()
+
+    if args.profile:
+        pr = cProfile.Profile()
+        pr.enable()
+        main(args.source, args.output, args.split)
+        pr.disable()
+        buf = io.StringIO()
+        ps = pstats.Stats(pr, stream=buf).sort_stats(pstats.SortKey.CUMULATIVE)
+        ps.print_stats(args.profile)
+        print("\n--- cProfile output ---")
+        print(buf.getvalue())
+    else:
+        main(args.source, args.output, args.split)

bae_chat_log_analysis/reducer.py

@@ -3,17 +3,75 @@ from collections import defaultdict
 from datetime import date
 from pathlib import Path
 
+import ahocorasick
 import emoji
 
+_NAME_REPLACEMENTS = {
+    'Honeybelle Ong-Jimenez Gaitano': 'Belle',
+    'Jason Ross': 'Jason',
+}
 
-def _reduce_line(line, abbreviations):
-    """Apply all token reductions to a single line."""
-    line = line.replace('Honeybelle Ong-Jimenez Gaitano', 'Belle')
-    line = line.replace('Jason Ross', 'Jason')
-    line = emoji.demojize(line)
-    for long, short in abbreviations.items():
-        line = line.replace(long, short)
-    return line
+
+def _build_automaton(abbreviations):
+    """Build an Aho-Corasick automaton for name and multi-codepoint emoji replacements.
+
+    Single-codepoint emoji are handled separately via str.translate() (_build_translate_table),
+    which is a pure-C O(n) pass and much faster for high-frequency single chars.
+    Abbreviation strings (':long_name:' -> ':short:') are omitted — they were only needed
+    as a second pass after emoji.demojize(), which we no longer call.
+    """
+    A = ahocorasick.Automaton()
+    for long, short in _NAME_REPLACEMENTS.items():
+        A.add_word(long, (len(long), short))
+    for char, meta in emoji.EMOJI_DATA.items():
+        if len(char) > 1:
+            en_name = meta['en']
+            target = abbreviations.get(en_name, en_name)
+            A.add_word(char, (len(char), target))
+    if len(A) == 0:
+        return None
+    A.make_automaton()
+    return A
+
+
+def _build_translate_table(abbreviations):
+    """Build a str.translate() table for single-codepoint emoji.
+
+    Must be applied AFTER the automaton pass so that multi-codepoint emoji (e.g. 👋🏻)
+    are consumed first, leaving only standalone single-codepoint occurrences.
+    """
+    table = {}
+    for char, meta in emoji.EMOJI_DATA.items():
+        if len(char) == 1:
+            en_name = meta['en']
+            table[ord(char)] = abbreviations.get(en_name, en_name)
+    return table
+
+
+def _apply_automaton(text, automaton):
+    """Single-pass replacement using Aho-Corasick (left-to-right, longest match wins)."""
+    # Keep only the longest match per start position to avoid sorting dominated matches.
+    # Composite emoji (e.g. 👋🏻) generate overlapping matches at the same start;
+    # the dict ensures only the longest is retained before we sort start positions.
+    best = {}
+    for end_idx, (pat_len, replacement) in automaton.iter(text):
+        start = end_idx - pat_len + 1
+        end = end_idx + 1
+        existing = best.get(start)
+        if existing is None or end > existing[0]:
+            best[start] = (end, replacement)
+
+    parts = []
+    last_end = 0
+    for start in sorted(best):
+        end, replacement = best[start]
+        if start < last_end:
+            continue  # overlapping with a previously committed match — skip
+        parts.append(text[last_end:start])
+        parts.append(replacement)
+        last_end = end
+    parts.append(text[last_end:])
+    return ''.join(parts)
 
 
 _DATE_RE = re.compile(r'^(\d{1,2})/(\d{1,2})/(\d{2,4}),')
@@ -32,37 +90,53 @@ def _date_key(month, day, year_short, split_by):
         return f"{year}-{int(month):02d}-{int(day):02d}"
 
 
-def reduce_tokens(input_file, output_file, abbreviations, split_by='none'):
+def _reduce_line(line, automaton, translate_table):
+    if automaton is not None:
+        line = _apply_automaton(line, automaton)
+    return line.translate(translate_table)
+
+
+def reduce_tokens(input_file, output_file, abbreviations, split_by='none', encoding=None):
     """Replace long usernames and convert emojis to shortened text codes.
 
     split_by: 'none', 'year', 'month', 'week', or 'day'
-    Returns a list of output file paths that were written.
+    encoding: optional tiktoken encoding; if provided, token counts are computed with
+              two encode calls total (source + all reduced) — no re-reads needed.
+    Returns (list[Path], int, int) — output paths, total reduced token count,
+    and source token count. Token counts are -1 when encoding is None.
     """
     output_file = Path(output_file)
     if output_file.suffix != '.txt':
         output_file = output_file.with_suffix(output_file.suffix + '.txt')
 
+    automaton = _build_automaton(abbreviations)
+    translate_table = _build_translate_table(abbreviations)
+
+    with open(input_file, 'r', encoding='utf-8') as f_in:
+        source_text = f_in.read()
+
+    source_tokens = len(encoding.encode(source_text)) if encoding is not None else -1
+
     if split_by == 'none':
-        with open(input_file, 'r', encoding='utf-8') as f_in, \
-             open(output_file, 'w', encoding='utf-8') as f_out:
-            for line in f_in:
-                f_out.write(_reduce_line(line, abbreviations))
-        return [Path(output_file)]
+        reduced = _reduce_line(source_text, automaton, translate_table)
+        with open(output_file, 'w', encoding='utf-8') as f_out:
+            f_out.write(reduced)
+        reduced_tokens = len(encoding.encode(reduced)) if encoding is not None else -1
+        return [Path(output_file)], reduced_tokens, source_tokens
 
     # Split mode: group lines by date bucket
     buckets = defaultdict(list)
     current_key = None
 
-    with open(input_file, 'r', encoding='utf-8') as f_in:
-        for line in f_in:
-            match = _DATE_RE.match(line)
-            if match:
-                month, day, year_short = match.groups()
-                current_key = _date_key(month, day, year_short, split_by)
-            reduced = _reduce_line(line, abbreviations)
-            if current_key is None:
-                current_key = '_unknown'
-            buckets[current_key].append(reduced)
+    for line in source_text.splitlines(keepends=True):
+        reduced = _reduce_line(line, automaton, translate_table)
+        match = _DATE_RE.match(reduced)
+        if match:
+            month, day, year_short = match.groups()
+            current_key = _date_key(month, day, year_short, split_by)
+        if current_key is None:
+            current_key = '_unknown'
+        buckets[current_key].append(reduced)
 
     output_path = Path(output_file)
     stem = output_path.stem if output_path.suffix else output_path.name
@@ -71,9 +145,13 @@ def reduce_tokens(input_file, output_file, abbreviations, split_by='none'):
 
     written_files = []
     for key in sorted(buckets):
+        content = ''.join(buckets[key])
         out = parent / f"{stem}-{key}{suffix}"
         with open(out, 'w', encoding='utf-8') as f_out:
-            f_out.writelines(buckets[key])
+            f_out.write(content)
         written_files.append(out)
 
-    return written_files
+    all_reduced = ''.join(''.join(lines) for lines in buckets.values())
+    reduced_tokens = len(encoding.encode(all_reduced)) if encoding is not None else -1
+
+    return written_files, reduced_tokens, source_tokens

requirements.txt

@@ -1,2 +1,3 @@
 emoji>=2.0.0
+pyahocorasick>=2.0.0
 tiktoken>=0.5.0