import numpy as np from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity from nltk.tokenize import sent_tokenize, word_tokenize from nltk.stem import WordNetLemmatizer from keybert import KeyBERT import re from collections import defaultdict from sklearn.cluster import AgglomerativeClustering class TopicSegmenter: def __init__(self, window_size=3, similarity_threshold=0.2, context_size=2, min_segment_size=2, topic_similarity_threshold=0.3, max_topics=8, hierarchical_threshold=0.6, model_name="all-MiniLM-L6-v2"): self.window_size = window_size self.similarity_threshold = similarity_threshold self.context_size = context_size self.min_segment_size = min_segment_size self.topic_similarity_threshold = topic_similarity_threshold self.max_topics = max_topics # Maximum number of main topics self.hierarchical_threshold = hierarchical_threshold # Threshold for hierarchical clustering try: # Explicitly set the model to a lightweight and compatible one import os os.environ["TOKENIZERS_PARALLELISM"] = "false" # Prevent parallelism warning # Try using sentence-transformers directly first try: from sentence_transformers import SentenceTransformer embedding_model = SentenceTransformer(model_name) self.keyword_extractor = KeyBERT(model=embedding_model) except (ImportError, ValueError, RuntimeError) as e: # Fallback to default model with lower-level settings self.keyword_extractor = KeyBERT(model='distilbert-base-nli-mean-tokens') except Exception as e: import logging logging.error(f"Error initializing KeyBERT model: {str(e)}") # Last resort fallback to simple keyword extraction without a model self.keyword_extractor = None self.vectorizer = TfidfVectorizer( stop_words='english', min_df=1, max_df=0.95, ngram_range=(1, 2) ) self.segment_vectorizer = TfidfVectorizer( stop_words='english', min_df=1, max_df=1.0, ngram_range=(1, 2) ) self.lemmatizer = WordNetLemmatizer() self.topic_history = [] self.topic_hierarchy = {} # Store parent-child relationships self.fitted_vectorizer = None def preprocess_text(self, text): # First, split the text by timestamps timestamp_chunks = re.findall(r'\[(\d+)\](.*?)(?=\[\d+\]|$)', text, re.DOTALL) all_sentences = [] original_sentences = [] timestamp_mapping = [] # To track which timestamp each sentence belongs to for timestamp, chunk_text in timestamp_chunks: # Split each chunk into sentences chunk_sentences = sent_tokenize(chunk_text.strip()) for sentence in chunk_sentences: if sentence.strip(): # Keep track of the original sentence original_sentences.append(sentence.strip()) # Clean and process the sentence for analysis clean_sentence = re.sub(r'^[^:]+:', '', sentence).strip() clean_sentence = re.sub(r'[^\w\s]', '', clean_sentence.lower()) words = word_tokenize(clean_sentence) lemmatized = [self.lemmatizer.lemmatize(word) for word in words] processed_sentence = ' '.join(lemmatized) all_sentences.append(processed_sentence) timestamp_mapping.append(int(timestamp)) return all_sentences, original_sentences, timestamp_mapping def get_topic_fingerprint(self, segment_text): cleaned_text = ' '.join([re.sub(r'^[^:]+:', '', sent).strip() for sent in segment_text]) if self.fitted_vectorizer is None: self.fitted_vectorizer = self.segment_vectorizer.fit([cleaned_text]) tfidf_matrix = self.fitted_vectorizer.transform([cleaned_text]) else: tfidf_matrix = self.fitted_vectorizer.transform([cleaned_text]) return tfidf_matrix.toarray()[0] def compare_with_previous_topics(self, current_segment): if not self.topic_history: return None, 0.0 current_fingerprint = self.get_topic_fingerprint(current_segment) max_similarity = 0.0 best_match_idx = None for idx, (topic_fingerprint, _, _, _, _) in enumerate(self.topic_history): if len(topic_fingerprint) != len(current_fingerprint): continue similarity = cosine_similarity([topic_fingerprint], [current_fingerprint])[0][0] if similarity > max_similarity: max_similarity = similarity best_match_idx = idx return best_match_idx, max_similarity def extract_keywords(self, sentences, top_n=3): if isinstance(sentences, list): text = ' '.join(sentences) else: text = sentences # If KeyBERT was not properly initialized, use TF-IDF as a fallback if self.keyword_extractor is None: try: # Use scikit-learn's TF-IDF for keyword extraction as fallback from sklearn.feature_extraction.text import TfidfVectorizer # Create a temporary vectorizer for this text only temp_vectorizer = TfidfVectorizer( max_features=50, stop_words='english', ngram_range=(1, 2) ) # Fit and transform on this specific text tfidf_matrix = temp_vectorizer.fit_transform([text]) feature_names = temp_vectorizer.get_feature_names_out() # Get top keywords based on TF-IDF scores tfidf_scores = tfidf_matrix.toarray()[0] scored_tokens = [(feature_names[i], tfidf_scores[i]) for i in range(len(feature_names))] sorted_tokens = sorted(scored_tokens, key=lambda x: x[1], reverse=True) return [token for token, score in sorted_tokens[:top_n]] except Exception as e: import logging logging.error(f"Error in TF-IDF fallback keyword extraction: {str(e)}") # If everything fails, just return some generic keywords return ["topic", "section", "content"] try: keywords = self.keyword_extractor.extract_keywords( text, keyphrase_ngram_range=(1, 2), # Allow for 1-2 word keyphrases stop_words='english', top_n=top_n ) sorted_keywords = sorted(keywords, key=lambda x: x[1], reverse=True) return [kw[0] for kw in sorted_keywords[:top_n]] except Exception as e: import logging logging.error(f"Error in KeyBERT keyword extraction: {str(e)}") # Fall back to TF-IDF method return self.extract_keywords(text, top_n) def detect_speaker_changes(self, text): """ Detect points in the transcript where speakers change. Returns a list of indices where speaker changes occur. """ speaker_pattern = re.compile(r'(?:^|\s)(?:I\'m|And I\'m|Hey I\'m|I am)\s+([A-Z][a-z]+)', re.MULTILINE) matches = list(speaker_pattern.finditer(text)) speaker_boundaries = [] for match in matches: # Find the sentence containing this speaker introduction sentences = sent_tokenize(text[:match.end()]) if sentences: speaker_boundaries.append(len(sentences) - 1) return speaker_boundaries def find_closest_timestamp(self, sentence, transcript): """ Find the timestamp after the most recent timestamp before the sentence. """ chunks = re.findall(r'\[(\d+)\](.*?)(?=\[\d+\]|$)', transcript, re.DOTALL) sentence = sentence.strip().lower() previous_timestamp = 0 current_timestamp = 0 for i, (timestamp, text) in enumerate(chunks): current_timestamp = int(timestamp) text = text.strip().lower() if sentence in text or text in sentence: if i > 0 and i < len(chunks): return current_timestamp return previous_timestamp previous_timestamp = current_timestamp return current_timestamp def calculate_topic_importance(self, segment, total_segments): """ Calculate importance score for a topic based on: 1. Length of the segment 2. Position in the transcript 3. Keyword significance """ # Length factor - longer segments likely more important length_score = min(1.0, len(segment) / 10) # Cap at 1.0 for segments of 10+ sentences # Position factor - beginning and end often contain important information position_in_doc = self.topic_history.index(segment) / total_segments position_score = 1.0 - min(abs(position_in_doc - 0.0), abs(position_in_doc - 1.0)) # Keyword significance - look for key terms indicating importance importance_indicators = ['important', 'key', 'main', 'primary', 'critical', 'essential', 'significant', 'revolutionary', 'breakthrough', 'innovative'] text = ' '.join(segment) significance_score = 0.0 for indicator in importance_indicators: if indicator in text.lower(): significance_score += 0.2 # Add 0.2 for each indicator found significance_score = min(1.0, significance_score) # Cap at 1.0 # Combined score, weighted combined_score = (0.4 * length_score) + (0.3 * position_score) + (0.3 * significance_score) return combined_score def cluster_topics(self, segments): """ Cluster similar topics together to form a hierarchy """ if len(segments) <= 1: return segments, [] # Extract features for each segment feature_vectors = [] for segment_idx, segment in enumerate(segments): feature_vectors.append(self.get_topic_fingerprint(segment)) feature_matrix = np.array(feature_vectors) # Use hierarchical clustering with compatible parameters num_clusters = min(self.max_topics, len(segments)) # Try different parameter combinations based on scikit-learn version compatibility try: # First attempt with basic parameters clustering = AgglomerativeClustering( n_clusters=num_clusters ) labels = clustering.fit_predict(feature_matrix) except Exception as e: print(f"Clustering with default parameters failed: {str(e)}") # Fallback to manual clustering based on similarity labels = self.manual_clustering(feature_matrix, num_clusters) # Group segments by cluster clusters = defaultdict(list) for segment_idx, cluster_id in enumerate(labels): clusters[cluster_id].append(segment_idx) # Create parent-child relationships parent_child_map = {} for cluster_id, segment_indices in clusters.items(): # If cluster has multiple segments, create a parent topic if len(segment_indices) > 1: # Combine segments to create a parent topic all_content = [] for idx in segment_indices: all_content.extend(segments[idx]) # Extract the most representative keywords keywords = self.extract_keywords(all_content, top_n=3) parent_topic = f"{', '.join(keywords)}" # Add parent-child relationships for idx in segment_indices: parent_child_map[idx] = parent_topic return labels, parent_child_map def manual_clustering(self, feature_matrix, num_clusters): """ Implement a simple similarity-based clustering as fallback """ n_samples = feature_matrix.shape[0] # If we have very few samples, each gets its own cluster if n_samples <= num_clusters: return np.arange(n_samples) # Calculate similarity matrix similarity = cosine_similarity(feature_matrix) labels = np.arange(n_samples) current_num_clusters = n_samples while current_num_clusters > num_clusters: max_similarity = -1 merge_i, merge_j = 0, 0 for i in range(n_samples): for j in range(i+1, n_samples): if labels[i] == labels[j]: continue if similarity[i, j] > max_similarity: max_similarity = similarity[i, j] merge_i, merge_j = i, j old_label = labels[merge_j] new_label = labels[merge_i] for i in range(n_samples): if labels[i] == old_label: labels[i] = new_label unique_labels = np.unique(labels) mapping = {old: new for new, old in enumerate(unique_labels)} labels = np.array([mapping[label] for label in labels]) current_num_clusters -= 1 return labels def segment_transcript(self, text): cleaned_sentences, original_sentences, timestamp_mapping = self.preprocess_text(text) if len(cleaned_sentences) == 0: return [], [], [], {} similarity_matrix = self.calculate_similarity_matrix(cleaned_sentences) content_boundaries = self.detect_topic_boundaries(similarity_matrix) speaker_boundaries = self.detect_speaker_changes(text) timestamp_boundaries = [] for i in range(1, len(timestamp_mapping)): if timestamp_mapping[i] - timestamp_mapping[i-1] > 5: # If there's a gap of more than 5 seconds timestamp_boundaries.append(i) all_boundaries = sorted(set(content_boundaries + speaker_boundaries + timestamp_boundaries)) initial_segments = [] initial_segment_texts = [] current_topic_id = 0 self.topic_history = [] # Reset topic history start_idx = 0 for boundary in all_boundaries + [len(original_sentences)]: if boundary - start_idx < self.min_segment_size: continue current_segment = original_sentences[start_idx:boundary] segment_timestamps = timestamp_mapping[start_idx:boundary] closest_timestamp = segment_timestamps[0] if segment_timestamps else 0 # Extract keywords directly from content without predefined categories keywords = self.extract_keywords(current_segment, top_n=4) # Create topic name directly from the top keywords - completely dynamic topic_name = f"{', '.join(keywords[:2])}" # Calculate importance score topic_importance = len(current_segment) / 3 # Simple approach: longer segments are more important # Create a topic fingerprint topic_fingerprint = self.get_topic_fingerprint(current_segment) # Store topic information self.topic_history.append((topic_fingerprint, topic_name, [current_segment], closest_timestamp, topic_importance)) initial_segments.append(current_segment) initial_segment_texts.append(topic_name) start_idx = boundary # Cluster segments to form hierarchical topics if len(initial_segments) > 1: cluster_labels, parent_child_map = self.cluster_topics(initial_segments) else: cluster_labels = [0] if initial_segments else [] parent_child_map = {} # Organize segments into hierarchical structure final_segments = [] topic_mappings = [] topic_hierarchies = {} # Group by parent topics parent_topic_groups = defaultdict(list) for i, segment in enumerate(initial_segments): topic_id = i # Check if this segment belongs to a parent topic if i in parent_child_map: parent_topic = parent_child_map[i] parent_id = len(self.topic_history) # If this is the first time seeing this parent topic if parent_topic not in parent_topic_groups: # Add the parent topic to our history all_content = [] for j, seg in enumerate(initial_segments): if j in parent_child_map and parent_child_map[j] == parent_topic: all_content.extend(seg) parent_fingerprint = self.get_topic_fingerprint(all_content) parent_timestamp = self.topic_history[i][3] # Use timestamp of first segment parent_importance = 1.0 # Parent topics are most important self.topic_history.append((parent_fingerprint, parent_topic, [all_content], parent_timestamp, parent_importance)) topic_hierarchies[parent_id] = [j for j, seg in enumerate(initial_segments) if j in parent_child_map and parent_child_map[j] == parent_topic] # Find the parent ID for j, (_, name, _, _, _) in enumerate(self.topic_history): if name == parent_topic: parent_id = j break # Add to parent topic group parent_topic_groups[parent_topic].append((segment, topic_id)) else: # This is already a main topic final_segments.append(segment) topic_mappings.append(topic_id) # Add grouped segments under their parent topics for parent_topic, segments_and_ids in parent_topic_groups.items(): # Find the parent ID parent_id = None for i, (_, name, _, _, _) in enumerate(self.topic_history): if name == parent_topic: parent_id = i break if parent_id is not None: # Add parent as a segment parent_content = [] for segment, _ in segments_and_ids: parent_content.extend(segment) final_segments.append(parent_content) topic_mappings.append(parent_id) # Add children segments for segment, child_id in segments_and_ids: final_segments.append(segment) topic_mappings.append(child_id) # For topics that don't have many segments, merge them if they're similar if len(final_segments) > self.max_topics * 2: merged_segments = [] merged_mappings = [] i = 0 while i < len(final_segments): current_segment = final_segments[i] current_mapping = topic_mappings[i] # If this is a short segment, try to merge with next if len(current_segment) <= 2 and i + 1 < len(final_segments): next_segment = final_segments[i+1] combined = current_segment + next_segment # Use keywords to see if they're related current_keywords = self.extract_keywords(current_segment) next_keywords = self.extract_keywords(next_segment) # Check for keyword overlap overlap = any(kw in next_keywords for kw in current_keywords) if overlap: # Merge the segments merged_segments.append(combined) merged_mappings.append(current_mapping) # Keep the first mapping i += 2 # Skip both segments continue # If no merge, keep as is merged_segments.append(current_segment) merged_mappings.append(current_mapping) i += 1 final_segments = merged_segments topic_mappings = merged_mappings return final_segments, topic_mappings, self.topic_history, topic_hierarchies def calculate_similarity_matrix(self, sentences): if not sentences: return np.array([[]]) tfidf_matrix = self.vectorizer.fit_transform(sentences) return cosine_similarity(tfidf_matrix) def detect_topic_boundaries(self, similarity_matrix): boundaries = [] n_sentences = len(similarity_matrix) if n_sentences <= 2 * self.window_size: return [] for i in range(self.window_size, n_sentences - self.window_size): prev_window = similarity_matrix[i-self.window_size:i, i-self.window_size:i] prev_similarity = np.mean(prev_window) next_window = similarity_matrix[i:i+self.window_size, i:i+self.window_size] next_similarity = np.mean(next_window) cross_window = similarity_matrix[i-self.window_size:i, i:i+self.window_size] cross_similarity = np.mean(cross_window) # Stricter conditions for detecting boundaries if (cross_similarity < self.similarity_threshold and cross_similarity < prev_similarity * 0.7 and # Increased constraint cross_similarity < next_similarity * 0.7 and # Increased constraint (len(boundaries) == 0 or i - boundaries[-1] >= self.min_segment_size)): boundaries.append(i) return boundaries if __name__ == "__main__": transcript = """ [0] Hey I'm Stanley, I'm a Stanford CS major. I did fun at Engineering at Facebook. [4] Hey I'm Andy, I'm also a Stanford CS major and I did platform engineering at Facebook. [8] I'm Evan, I was on the founding team of Vivo, the music video service. [12] And I'm Tony, I was a product editor at Square. [15] And the four of us came together about six months ago to work on software for small business owners. [20] But we didn't have a need at first so we just went out and talked to all the small business owners we could find. [24] [6] After over a hundred interviews, we came across a really interesting problem with small business restaurants in an area like this. [30] Yeah, so it turns out restaurants in Palo Alto don't deliver even though they really [34] want to, but they can't afford it. [36] But their consumers are craving for it, but the places that the consumers love just can't [41] deliver. [42] And we also found out about these delivery drivers who had a ton of spare time and they [45] all wanted an extra cash during that downtime. [48] Right, so that's when we built an initial product, PaloAltodelivery.com, and how it [53] worked is the customer goes to the website and places it all in order that gets automatically [59] [32] since the restaurant [60] And then we as the dispatchers with some pretty neat routing and badging algorithms were able to send the drivers efficiently to [66] Get the orders to the customers at a really fast time. Yeah, and the four of us actually started off as delivery drivers and [74] Over time we hired more as we grew and in our first month of launch with not much marketing [80] We got over 150 paying customers in the Palo Alto area, which was really awesome [84] And from that we generated over $10,000 in sales """ segmenter = TopicSegmenter( window_size=2, # Smaller window for more granular segmentation similarity_threshold=0.15, # Lower threshold to detect more subtle changes context_size=1, min_segment_size=2, # Smaller minimum size for segments topic_similarity_threshold=0.25, # Lower threshold to better distinguish topics max_topics=5, # Target number of main topics hierarchical_threshold=0.6 # Threshold for hierarchical clustering ) segments, topic_mappings, topic_history, topic_hierarchies = segmenter.segment_transcript(transcript) print("Topic Segmentation Analysis:\n") print(f"Found {len(set(topic_mappings))} distinct topics") # Print hierarchical topic structure print("\nTopic Hierarchy:") for parent_id, child_ids in topic_hierarchies.items(): parent_name = topic_history[parent_id][1] print(f"Main Topic: {parent_name}") for child_id in child_ids: child_name = topic_history[child_id][1] print(f" - Subtopic: {child_name}") print("\nDetailed Segments:") for i, (segment, topic_id) in enumerate(zip(segments, topic_mappings)): print(f"Segment {i+1} (Part of {topic_history[topic_id][1]}):") print(f"Closest Timestamp: {topic_history[topic_id][3]} seconds") print("-" * 50) print("\n".join(segment[:2]) + "..." if len(segment) > 2 else "\n".join(segment)) print("-" * 50 + "\n")