Update README.md

README.md

@@ -56,9 +56,21 @@ pip install -U sentence-transformers
 Then you can implement like this:
 
 ```python
----
-license: mit
----
+"""
+Text Chunking Utility
+
+This module provides functionality to intelligently chunk text documents into semantically coherent sections
+using sentence embeddings and cosine similarity. It's particularly useful for processing large documents
+while maintaining contextual relationships between sentences.
+
+Requirements:
+    - nltk
+    - sentence-transformers
+    - scikit-learn
+    - numpy
+    - matplotlib
+"""
+
 import nltk
 from nltk.tokenize import sent_tokenize
 from sentence_transformers import SentenceTransformer
@@ -66,94 +78,138 @@ from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 import matplotlib.pyplot as plt
 
-# Load and tokenize the text
-def load_and_tokenize(file_path):
-    with open(file_path, 'r') as f:
-        text = f.read()
-    return sent_tokenize(text)
-
-# Combine sentences with their neighbors
-def combine_sentences(sentences, buffer=1):
-    combined = []
-    for i in range(len(sentences)):
-        start = max(0, i - buffer)
-        end = min(len(sentences), i + buffer + 1)
-        combined.append(' '.join(sentences[start:end]))
-    return combined
-
-# Calculate cosine distances between embeddings
-def calc_cosine_distances(embeddings):
-    distances = []
-    for i in range(len(embeddings) - 1):
-        sim = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0]
-        distances.append(1 - sim)
-    return distances
-
-# Find breakpoints based on distance threshold
-def find_breakpoints(distances, percentile=95):
-    threshold = np.percentile(distances, percentile)
-    return [i for i, d in enumerate(distances) if d > threshold]
-
-# Create chunks based on breakpoints
-def create_chunks(sentences, breakpoints):
-    chunks = []
-    start = 0
-    for bp in breakpoints:
-        chunks.append(' '.join(sentences[start:bp + 1]))
-        start = bp + 1
-    chunks.append(' '.join(sentences[start:]))
-    return chunks
-
-# Merge small chunks with their most similar neighbor
-def merge_small_chunks(chunks, embeddings, min_size=3):
-    merged = [chunks[0]]
-    merged_emb = [embeddings[0]]
-    
-    for i in range(1, len(chunks) - 1):
-        if len(chunks[i].split('. ')) < min_size:
-            prev_sim = cosine_similarity([embeddings[i]], [merged_emb[-1]])[0][0]
-            next_sim = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0]
+
+class TextChunker:
+    def __init__(self, model_name='sentence-transformers/all-mpnet-base-v1'):
+        """Initialize the TextChunker with a specified sentence transformer model."""
+        self.model = SentenceTransformer(model_name)
+
+    def process_file(self, file_path, context_window=1, percentile_threshold=95, min_chunk_size=3):
+        """
+        Process a text file and split it into semantically meaningful chunks.
+        
+        Args:
+            file_path: Path to the text file
+            context_window: Number of sentences to consider on either side for context
+            percentile_threshold: Percentile threshold for identifying breakpoints
+            min_chunk_size: Minimum number of sentences in a chunk
             
-            if prev_sim > next_sim:
-                merged[-1] += ' ' + chunks[i]
-                merged_emb[-1] = (merged_emb[-1] + embeddings[i]) / 2
+        Returns:
+            list: Semantically coherent text chunks
+        """
+        # Process the text file
+        sentences = self._load_text(file_path)
+        contextualized = self._add_context(sentences, context_window)
+        embeddings = self.model.encode(contextualized)
+        
+        # Create and refine chunks
+        distances = self._calculate_distances(embeddings)
+        breakpoints = self._identify_breakpoints(distances, percentile_threshold)
+        initial_chunks = self._create_chunks(sentences, breakpoints)
+        
+        # Merge small chunks for better coherence
+        chunk_embeddings = self.model.encode(initial_chunks)
+        final_chunks = self._merge_small_chunks(initial_chunks, chunk_embeddings, min_chunk_size)
+        
+        return final_chunks
+
+    def _load_text(self, file_path):
+        """Load and tokenize text from a file."""
+        with open(file_path, 'r', encoding='utf-8') as file:
+            text = file.read()
+        return sent_tokenize(text)
+
+    def _add_context(self, sentences, window_size):
+        """Combine sentences with their neighbors for better context."""
+        contextualized = []
+        for i in range(len(sentences)):
+            start = max(0, i - window_size)
+            end = min(len(sentences), i + window_size + 1)
+            context = ' '.join(sentences[start:end])
+            contextualized.append(context)
+        return contextualized
+
+    def _calculate_distances(self, embeddings):
+        """Calculate cosine distances between consecutive embeddings."""
+        distances = []
+        for i in range(len(embeddings) - 1):
+            similarity = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0]
+            distance = 1 - similarity
+            distances.append(distance)
+        return distances
+
+    def _identify_breakpoints(self, distances, threshold_percentile):
+        """Find natural breaking points in the text based on semantic distances."""
+        threshold = np.percentile(distances, threshold_percentile)
+        return [i for i, dist in enumerate(distances) if dist > threshold]
+
+    def _create_chunks(self, sentences, breakpoints):
+        """Create initial text chunks based on identified breakpoints."""
+        chunks = []
+        start_idx = 0
+        
+        for breakpoint in breakpoints:
+            chunk = ' '.join(sentences[start_idx:breakpoint + 1])
+            chunks.append(chunk)
+            start_idx = breakpoint + 1
+            
+        # Add the final chunk
+        final_chunk = ' '.join(sentences[start_idx:])
+        chunks.append(final_chunk)
+        
+        return chunks
+
+    def _merge_small_chunks(self, chunks, embeddings, min_size):
+        """Merge small chunks with their most similar neighbor."""
+        final_chunks = [chunks[0]]
+        merged_embeddings = [embeddings[0]]
+        
+        for i in range(1, len(chunks) - 1):
+            current_chunk_size = len(chunks[i].split('. '))
+            
+            if current_chunk_size < min_size:
+                # Calculate similarities
+                prev_similarity = cosine_similarity([embeddings[i]], [merged_embeddings[-1]])[0][0]
+                next_similarity = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0]
+                
+                if prev_similarity > next_similarity:
+                    # Merge with previous chunk
+                    final_chunks[-1] = f"{final_chunks[-1]} {chunks[i]}"
+                    merged_embeddings[-1] = (merged_embeddings[-1] + embeddings[i]) / 2
+                else:
+                    # Merge with next chunk
+                    chunks[i + 1] = f"{chunks[i]} {chunks[i + 1]}"
+                    embeddings[i + 1] = (embeddings[i] + embeddings[i + 1]) / 2
             else:
-                chunks[i + 1] = chunks[i] + ' ' + chunks[i + 1]
-                embeddings[i + 1] = (embeddings[i] + embeddings[i + 1]) / 2
-        else:
-            merged.append(chunks[i])
-            merged_emb.append(embeddings[i])
-    
-    merged.append(chunks[-1])
-    merged_emb.append(embeddings[-1])
-    return merged, merged_emb
-
-# Main process
-def chunk_text(file_path):
-    # Load the model
-    model = SentenceTransformer('sentence-transformers/all-mpnet-base-v1')
-    
-    # Process the text
-    sentences = load_and_tokenize(file_path)
-    combined = combine_sentences(sentences)
-    embeddings = model.encode(combined)
-    
-    # Find breakpoints and create initial chunks
-    distances = calc_cosine_distances(embeddings)
-    breakpoints = find_breakpoints(distances)
-    chunks = create_chunks(sentences, breakpoints)
+                final_chunks.append(chunks[i])
+                merged_embeddings.append(embeddings[i])
+        
+        final_chunks.append(chunks[-1])
+        return final_chunks
+
+
+def main():
+    """Example usage of the TextChunker class."""
+    # Initialize the chunker
+    chunker = TextChunker()
     
-    # Merge small chunks
-    chunk_embeddings = model.encode(chunks)
-    final_chunks, _ = merge_small_chunks(chunks, chunk_embeddings)
+    # Process a text file
+    file_path = "path/to/your/document.txt"
+    chunks = chunker.process_file(
+        file_path,
+        context_window=1,
+        percentile_threshold=95,
+        min_chunk_size=3
+    )
     
-    return final_chunks
+    # Print results
+    print(f"Successfully split text into {len(chunks)} chunks")
+    print("\nFirst chunk preview:")
+    print(f"{chunks[0][:200]}...")
+
 
 if __name__ == "__main__":
-    file_path = "/path/to/your/text/file.txt"
-    result = chunk_text(file_path)
-    print(f"Number of chunks: {len(result)}")
-    print("First chunk:", result[0][:100] + "...")  
+    main()
 ```
 ## Evaluation Results