name_clusters.py

"""
Name clusters using LLM to generate theme names and descriptions.

Usage:
    modal run insights_first/name_clusters.py --input insights_first/data/clusters_20260120_030333.json
"""

import modal
import json
import argparse
from datetime import datetime

app = modal.App("cluster-namer")

MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"
model_volume = modal.Volume.from_name("qwen-model-cache", create_if_missing=True)

vllm_image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("vllm>=0.6.0", "torch", "transformers", "huggingface_hub")
)

NAMING_PROMPT = """You are naming a thread of connected insights extracted from {num_episodes} DIFFERENT podcast episodes.

This is an "invisible thread" - a common theme that appears across multiple conversations without the speakers knowing they're saying the same thing.

Your job is to find the ACTUAL common thread - the specific concept that genuinely connects these insights. Be concrete and precise, not impressive-sounding.

## Source episodes (insights came from these {num_episodes} different conversations):
{episode_list}

RULES:
1. NO grandiose or fluffy language ("Unlocking", "Mastering", "The Power of", "Transformative")
2. NO vague categories ("Leadership Wisdom", "Growth Strategies", "Building Culture")
3. Name should be LITERAL - someone reading it should know exactly what the thread is about
4. If insights don't have a clear common thread, say so
5. The name MUST reflect what the insights ACTUALLY say, not what you think they should say

GOOD thread names (specific, concrete):
- "PLG prioritizes user growth over revenue"
- "Fire underperformers within 30 days"
- "Feedback should be 90% shorter"
- "Freemium builds market share first"

BAD thread names (vague, fluffy, or mismatched):
- "The Art of Effective Leadership"
- "Building High-Performance Teams"
- "Remote Team Communication" (when insights are about something else)

## Insights in this thread:
{insights}

## Respond with:
THREAD_NAME: [The literal common thread in 3-7 words - be boringly specific]
CORE_CLAIM: [The single claim these insights support, in one plain sentence]
WHY_CONNECTED: [In 10 words or less, what actually connects these insights?]

If the insights don't have a clear common thread, respond:
THREAD_NAME: NO_CLEAR_THREAD
CORE_CLAIM: These insights are too diverse to form a coherent thread.
WHY_CONNECTED: N/A"""


@app.cls(
    gpu="A10G",
    image=vllm_image,
    volumes={"/model-cache": model_volume},
    timeout=600,
    scaledown_window=300,
)
class ClusterNamer:
    @modal.enter()
    def load_model(self):
        from vllm import LLM, SamplingParams
        
        self.llm = LLM(
            model=MODEL_ID,
            download_dir="/model-cache",
            trust_remote_code=True,
            max_model_len=4096,
            gpu_memory_utilization=0.9,
        )
        self.sampling_params = SamplingParams(
            temperature=0.7,  # Slightly creative
            max_tokens=300,
        )
    
    @modal.method()
    def name_cluster(self, cluster_name: str, insights: list[str], num_episodes: int = 3, episode_list: str = "") -> dict:
        """Name a single cluster based on its insights."""
        # Use top 10 insights to avoid context overflow
        insights_text = "\n".join(f"- {ins}" for ins in insights[:10])
        prompt = NAMING_PROMPT.format(
            insights=insights_text,
            num_episodes=num_episodes,
            episode_list=episode_list or "(episodes not listed)"
        )
        
        outputs = self.llm.generate([prompt], self.sampling_params)
        response = outputs[0].outputs[0].text.strip()
        
        # Parse response
        result = {
            "original_name": cluster_name,
            "thread_name": "",
            "core_claim": "",
            "why_connected": "",
            "raw_response": response,
        }
        
        for line in response.split("\n"):
            line = line.strip()
            if line.upper().startswith("THREAD_NAME:"):
                result["thread_name"] = line[12:].strip()
            elif line.upper().startswith("CORE_CLAIM:"):
                result["core_claim"] = line[11:].strip()
            elif line.upper().startswith("WHY_CONNECTED:"):
                result["why_connected"] = line[14:].strip()
        
        return result


@app.local_entrypoint()
def main(input: str, output: str = None, min_episodes: int = 3):
    """Name all clusters in the input file."""
    
    # Load threads
    print(f"Loading threads from {input}...")
    with open(input, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    # Handle both cluster format and thread format
    if 'threads' in data:
        items = data['threads']
        # Thread format - list of thread objects
        if isinstance(items, list):
            data['clusters'] = {f"Thread {t['thread_id']}": t for t in items}
        else:
            data['clusters'] = items
    
    clusters = data['clusters']
    print(f"Found {len(clusters)} threads to process")
    
    # Pre-filter by min_episodes (true invisible threads must span multiple episodes)
    valid_clusters = {}
    skipped_single_source = 0
    for name, summary in clusters.items():
        if name == "Unclustered (Noise)":
            continue
        
        num_episodes = summary.get('num_episodes', 1)
        if num_episodes < min_episodes:
            skipped_single_source += 1
            print(f"  SKIP: {name} (only {num_episodes} episode(s))")
            continue
        
        valid_clusters[name] = summary
    
    print(f"Filtered to {len(valid_clusters)} threads with >={min_episodes} episodes")
    if skipped_single_source > 0:
        print(f"  Skipped {skipped_single_source} single-source threads")
    
    # Prepare inputs with episode context
    cluster_inputs = []
    for name, summary in valid_clusters.items():
        # Get insight texts
        if 'top_insights' in summary:
            insight_texts = summary['top_insights']
        elif 'insights' in summary:
            insight_texts = [i['insight_text'] for i in summary['insights'][:10]]
        else:
            continue
        
        # Get episode list for context
        episodes = summary.get('episodes', [])
        num_episodes = len(episodes)
        episode_list = "\n".join(f"- {ep[:60]}..." for ep in episodes[:5])
        if len(episodes) > 5:
            episode_list += f"\n- ...and {len(episodes) - 5} more"
        
        cluster_inputs.append((name, insight_texts, num_episodes, episode_list))
    
    # Name clusters in parallel
    namer = ClusterNamer()
    
    import time
    start = time.time()
    
    named_threads = {}
    for (cluster_name, insights, num_eps, ep_list), result in zip(
        cluster_inputs, 
        namer.name_cluster.starmap(cluster_inputs)
    ):
        named_threads[cluster_name] = {
            **valid_clusters[cluster_name],
            'thread_name': result['thread_name'],
            'core_claim': result['core_claim'],
            'why_connected': result['why_connected'],
        }
        print(f"  {cluster_name} ({num_eps} eps) → \"{result['thread_name']}\"")
    
    elapsed = time.time() - start
    print(f"\nNaming complete in {elapsed:.1f}s")
    
    # Save output
    if output is None:
        output = input.replace("clusters_", "named_threads_").replace("threads_", "named_threads_")
    
    output_data = {
        'metadata': {
            **data['metadata'],
            'naming_timestamp': datetime.now().isoformat(),
            'naming_model': MODEL_ID,
            'min_episodes_filter': min_episodes,
            'skipped_single_source': skipped_single_source,
        },
        'threads': named_threads,
    }
    
    with open(output, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, indent=2, ensure_ascii=False)
    
    print(f"\nResults saved to: {output}")
    
    # Print summary
    print("\n" + "="*60)
    print("NAMED THREADS")
    print("="*60)
    
    # Sort by size
    sorted_threads = sorted(
        named_threads.items(),
        key=lambda x: x[1]['size'],
        reverse=True
    )
    
    for name, thread in sorted_threads:
        if thread['thread_name'] == 'NO_CLEAR_THREAD':
            print(f"\n❌ {name}: NO CLEAR THREAD ({thread['size']} insights)")
        else:
            print(f"\n🧵 {thread['thread_name']} ({thread['size']} insights, {thread['num_episodes']} episodes)")
            print(f"   Core claim: {thread['core_claim']}")