Update README.md

README.md

@@ -66,90 +66,94 @@ from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 import matplotlib.pyplot as plt
 
-# Text to be chunked
-with open("/path to text") as f:
-  text = f.read()
-
-# Tokenize the text into sentences
-sentences = sent_tokenize(text)
-
-# Generate embeddings for each sentence using sentence-transformers model of choice
-model = SentenceTransformer('sentence-transformers/all-mpnet-base-v1')
-embeddings = model.encode(sentences)
-
-# Combine the sentences with their neighbors
-# Adjust buffer size to change how many neighboring sentences on either side of a target sentence are included in the combined text (1=1 before and after). 
-def combine_sentences(sentences, buffer_size=1):
-  combined_sentences = []
-  for i in range(len(sentences)):
-    combined_sentence = ' '.join(sentences[max(0, i-buffer_size):min(len(sentences), i+1+buffer_size)])
-    combined_sentences.append(combined_sentence)
-  return combined_sentences
-
-combined_sentences = combine_sentences(sentences)
-combined_embeddings = model.encode(combined_sentences)
+# 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 calculate_cosine_distances(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
-
-distances = calculate_cosine_distances(combined_embeddings)
-
-# Identify breakpoints
-# Adjust breakpoint threshhold to change the level of dissimilarity between chunk embeddings (higher for greater dissimilarity)
-breakpoint_percentile_threshold = 95
-breakpoint_distance_threshold = np.percentile(distances, breakpoint_percentile_threshold)
-breakpoint_indices = [i for i, distance in enumerate(distances) if distance > breakpoint_distance_threshold]
+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
-chunks = []
-start_index = 0
-for breakpoint_index in breakpoint_indices:
-  chunk = ' '.join(sentences[start_index:breakpoint_index + 1])
-  chunks.append(chunk)
-  start_index = breakpoint_index + 1
-chunks.append(' '.join(sentences[start_index:]))
-
-# Set a minimum number of sentences per chunk
-min_chunk_size = 3
-
-# Merge small chunks with their most semantically similar neighbor
-def merge_small_chunks_with_neighbors(chunks, embeddings):
-  merged_chunks = [chunks[0]] # Start with the first chunk
-  merged_embeddings = [embeddings[0]] # And its embedding
-
-  for i in range(1, len(chunks) - 1): # Iterate through chunks, excluding the first and last
-    # If the current chunk is small, consider merging it with a neighbor
-    if len(chunks[i].split('. ')) < min_chunk_size:
-      prev_similarity = cosine_similarity([embeddings[i]], [merged_embeddings[-1]])[0][0]
-      next_similarity = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0]
-
-      # Merge with the most similar neighbor
-      if prev_similarity > next_similarity:
-        merged_chunks[-1] += ' ' + chunks[i]
-        merged_embeddings[-1] = (merged_embeddings[-1] + embeddings[i]) / 2
-      else:
-        chunks[i + 1] = chunks[i] + ' ' + chunks[i + 1]
-        embeddings[i + 1] = (embeddings[i] + embeddings[i + 1]) / 2
-    else:
-      merged_chunks.append(chunks[i])
-      merged_embeddings.append(embeddings[i])
-
-  merged_chunks.append(chunks[-1])
-  merged_embeddings.append(embeddings[-1])
-
-  return merged_chunks, merged_embeddings
-
-# Generate embeddings for each initial chunk and merge most semantically similar neighbors
-chunk_embeddings = model.encode(chunks)
-chunks, chunk_embeddings = merge_small_chunks_with_neighbors(chunks, chunk_embeddings)
-
-print(chunks[0])
+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]
+            
+            if prev_sim > next_sim:
+                merged[-1] += ' ' + chunks[i]
+                merged_emb[-1] = (merged_emb[-1] + embeddings[i]) / 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)
+    
+    # Merge small chunks
+    chunk_embeddings = model.encode(chunks)
+    final_chunks, _ = merge_small_chunks(chunks, chunk_embeddings)
+    
+    return final_chunks
+
+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] + "...")  
 ```
 ## Evaluation Results