Spaces:
Runtime error
Runtime error
| from pathlib import Path | |
| import spaces | |
| import gradio as gr | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import json | |
| import random | |
| import re | |
| import time | |
| import uuid | |
| import zipfile | |
| from bs4 import BeautifulSoup | |
| from docx import Document as DocxDocument | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from pypdf import PdfReader | |
| from sentence_transformers import SentenceTransformer | |
| from transformers import AutoTokenizer | |
| SEED = 42 | |
| MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2" | |
| BATCH_SIZE = 8 | |
| DEVICE = "cpu" | |
| random.seed(SEED) | |
| np.random.seed(SEED) | |
| torch.manual_seed(SEED) | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) | |
| _embedding_model = None | |
| def get_embedding_model(): | |
| global _embedding_model | |
| if _embedding_model is None: | |
| _embedding_model = SentenceTransformer(MODEL_ID, device=DEVICE) | |
| return _embedding_model | |
| def clean_text(value): | |
| value = value.replace("\x00", " ") | |
| value = value.replace("\r\n", "\n").replace("\r", "\n") | |
| value = re.sub(r"[ \t]+", " ", value) | |
| value = re.sub(r" *\n *", "\n", value) | |
| value = re.sub(r"\n{3,}", "\n\n", value) | |
| return value.strip() | |
| def extract_pdf_page(page): | |
| try: | |
| return page.extract_text(extraction_mode="layout") or "" | |
| except TypeError: | |
| return page.extract_text() or "" | |
| def extract_document(path): | |
| path = Path(path) | |
| suffix = path.suffix.lower() | |
| page_count = None | |
| if suffix == ".pdf": | |
| reader = PdfReader(str(path)) | |
| pages = [extract_pdf_page(page) for page in reader.pages] | |
| raw_text = "\n\n".join(pages) | |
| page_count = len(pages) | |
| elif suffix == ".docx": | |
| document = DocxDocument(str(path)) | |
| raw_text = "\n\n".join( | |
| paragraph.text | |
| for paragraph in document.paragraphs | |
| if paragraph.text.strip() | |
| ) | |
| elif suffix == ".txt": | |
| raw_text = path.read_text(encoding="utf-8", errors="ignore") | |
| elif suffix in {".html", ".htm"}: | |
| html = path.read_text(encoding="utf-8", errors="ignore") | |
| soup = BeautifulSoup(html, "lxml") | |
| for tag in soup(["script", "style", "nav", "footer"]): | |
| tag.decompose() | |
| raw_text = soup.get_text("\n") | |
| else: | |
| raise ValueError("Upload a PDF, DOCX, TXT, HTML or HTM file.") | |
| return clean_text(raw_text), page_count, suffix[1:].upper() | |
| def token_count(value): | |
| return len( | |
| tokenizer.encode( | |
| value, | |
| add_special_tokens=False, | |
| truncation=False, | |
| verbose=False, | |
| ) | |
| ) | |
| def split_sentences(value): | |
| return [ | |
| sentence.strip() | |
| for sentence in re.split(r"(?<=[.!?])\s+|\n+", value) | |
| if sentence.strip() | |
| ] | |
| def fixed_chunks(value, chunk_size, overlap_ratio): | |
| token_ids = tokenizer.encode(value, add_special_tokens=False) | |
| overlap = int(chunk_size * overlap_ratio) | |
| step = max(1, chunk_size - overlap) | |
| chunks = [] | |
| for start in range(0, len(token_ids), step): | |
| piece = token_ids[start:start + chunk_size] | |
| if not piece: | |
| continue | |
| chunk = tokenizer.decode(piece, skip_special_tokens=True).strip() | |
| if chunk: | |
| chunks.append(chunk) | |
| if start + chunk_size >= len(token_ids): | |
| break | |
| return chunks | |
| def merge_units(units, chunk_size, overlap_ratio): | |
| chunks = [] | |
| current = [] | |
| for unit in units: | |
| if token_count(unit) > chunk_size: | |
| if current: | |
| chunks.append(" ".join(current).strip()) | |
| current = [] | |
| chunks.extend(fixed_chunks(unit, chunk_size, overlap_ratio)) | |
| continue | |
| candidate = " ".join(current + [unit]).strip() | |
| if current and token_count(candidate) > chunk_size: | |
| chunks.append(" ".join(current).strip()) | |
| overlap_budget = int(chunk_size * overlap_ratio) | |
| carried = [] | |
| for previous in reversed(current): | |
| test = " ".join([previous] + carried).strip() | |
| if token_count(test) <= overlap_budget: | |
| carried.insert(0, previous) | |
| else: | |
| break | |
| current = carried + [unit] | |
| else: | |
| current.append(unit) | |
| if current: | |
| chunks.append(" ".join(current).strip()) | |
| return [chunk for chunk in chunks if chunk] | |
| def sentence_chunks(value, chunk_size, overlap_ratio): | |
| return merge_units(split_sentences(value), chunk_size, overlap_ratio) | |
| def paragraph_chunks(value, chunk_size, overlap_ratio): | |
| paragraphs = [ | |
| paragraph.strip() | |
| for paragraph in re.split(r"\n\s*\n", value) | |
| if paragraph.strip() | |
| ] | |
| return merge_units(paragraphs, chunk_size, overlap_ratio) | |
| def recursive_chunks(value, chunk_size, overlap_ratio): | |
| splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer( | |
| tokenizer, | |
| chunk_size=chunk_size, | |
| chunk_overlap=int(chunk_size * overlap_ratio), | |
| separators=["\n\n", "\n", ". ", " ", ""], | |
| ) | |
| return [chunk.strip() for chunk in splitter.split_text(value) if chunk.strip()] | |
| def semantic_chunks(value, target_size, break_percentile=25): | |
| sentences = split_sentences(value) | |
| if len(sentences) <= 1: | |
| return sentences | |
| vectors = get_embedding_model().encode( | |
| sentences, | |
| batch_size=BATCH_SIZE, | |
| normalize_embeddings=True, | |
| convert_to_numpy=True, | |
| show_progress_bar=False, | |
| ) | |
| similarities = np.sum(vectors[:-1] * vectors[1:], axis=1) | |
| threshold = np.percentile(similarities, break_percentile) | |
| chunks = [] | |
| current = [sentences[0]] | |
| for index in range(1, len(sentences)): | |
| candidate = " ".join(current + [sentences[index]]) | |
| topic_shift = similarities[index - 1] < threshold | |
| too_large = token_count(candidate) > target_size | |
| if topic_shift or too_large: | |
| chunks.append(" ".join(current).strip()) | |
| current = [sentences[index]] | |
| else: | |
| current.append(sentences[index]) | |
| if current: | |
| chunks.append(" ".join(current).strip()) | |
| return chunks | |
| def create_chunks(text, strategy, chunk_size, overlap_percent): | |
| overlap_ratio = overlap_percent / 100 | |
| if strategy == "Fixed": | |
| return fixed_chunks(text, chunk_size, overlap_ratio) | |
| if strategy == "Sentence": | |
| return sentence_chunks(text, chunk_size, overlap_ratio) | |
| if strategy == "Paragraph": | |
| return paragraph_chunks(text, chunk_size, overlap_ratio) | |
| if strategy == "Recursive": | |
| return recursive_chunks(text, chunk_size, overlap_ratio) | |
| if strategy == "Semantic": | |
| return semantic_chunks(text, chunk_size) | |
| raise ValueError(f"Unsupported strategy: {strategy}") | |
| def boundary_score(chunks): | |
| if not chunks: | |
| return 0.0 | |
| valid = sum( | |
| bool(re.search(r"[.!?][\"')\]]?$", chunk.strip())) | |
| for chunk in chunks | |
| ) | |
| return valid / len(chunks) | |
| def context_preservation_score(source_text, chunks, sample_size=100): | |
| candidates = [ | |
| sentence | |
| for sentence in split_sentences(source_text) | |
| if token_count(sentence) >= 8 | |
| ] | |
| if not candidates: | |
| return 0.0 | |
| sample = random.Random(SEED).sample( | |
| candidates, | |
| min(sample_size, len(candidates)), | |
| ) | |
| preserved = sum(any(sentence in chunk for chunk in chunks) for sentence in sample) | |
| return preserved / len(sample) | |
| def run_benchmark(uploaded_file, strategy, chunk_size, overlap_percent, maximum_chunks): | |
| if uploaded_file is None: | |
| raise gr.Error("Upload one document first.") | |
| run_directory = Path("/tmp") / f"benchmark_{uuid.uuid4().hex[:8]}" | |
| run_directory.mkdir(parents=True, exist_ok=True) | |
| try: | |
| text, page_count, document_type = extract_document(uploaded_file) | |
| except Exception as error: | |
| raise gr.Error(f"Document extraction failed: {error}") | |
| if len(text) < 100: | |
| raise gr.Error("Very little text was extracted. A scanned PDF may require OCR.") | |
| document_summary = pd.DataFrame([{ | |
| "File": Path(uploaded_file).name, | |
| "Type": document_type, | |
| "Pages": page_count, | |
| "Characters": len(text), | |
| "Words": len(text.split()), | |
| "Tokens": token_count(text), | |
| }]) | |
| started = time.perf_counter() | |
| chunks = create_chunks(text, strategy, int(chunk_size), int(overlap_percent)) | |
| chunk_seconds = time.perf_counter() - started | |
| if not chunks: | |
| raise gr.Error("The selected configuration produced no chunks.") | |
| lengths = [token_count(chunk) for chunk in chunks] | |
| duplicates = len(chunks) - len(set(chunks)) | |
| effective_overlap = 0 if strategy == "Semantic" else int(overlap_percent) | |
| chunk_summary = pd.DataFrame([{ | |
| "Strategy": strategy, | |
| "Target tokens": int(chunk_size), | |
| "Overlap": f"{effective_overlap}%", | |
| "Chunks": len(chunks), | |
| "Average tokens": round(float(np.mean(lengths)), 2), | |
| "Minimum tokens": int(min(lengths)), | |
| "Maximum tokens": int(max(lengths)), | |
| "Boundary score": round(boundary_score(chunks), 4), | |
| "Context preservation": round(context_preservation_score(text, chunks), 4), | |
| "Duplicate ratio": round(duplicates / max(len(chunks), 1), 4), | |
| "Processing seconds": round(chunk_seconds, 4), | |
| }]) | |
| preview = "\n\n".join( | |
| f"CHUNK {index}\n{'=' * 60}\n{chunk[:1200]}" | |
| for index, chunk in enumerate(chunks[:3], start=1) | |
| ) | |
| selected_chunks = chunks[:int(maximum_chunks)] | |
| model = get_embedding_model() | |
| started = time.perf_counter() | |
| vectors = model.encode( | |
| selected_chunks, | |
| batch_size=BATCH_SIZE, | |
| normalize_embeddings=True, | |
| convert_to_numpy=True, | |
| show_progress_bar=False, | |
| ) | |
| embedding_seconds = time.perf_counter() - started | |
| test_pairs = { | |
| "Related": [ | |
| "Chunking divides a document into smaller retrieval units.", | |
| "Documents are split into smaller passages before retrieval.", | |
| ], | |
| "Unrelated": [ | |
| "Chunking divides a document into smaller retrieval units.", | |
| "The football match was postponed because of heavy rain.", | |
| ], | |
| } | |
| similarity_rows = [] | |
| for pair_name, pair in test_pairs.items(): | |
| pair_vectors = model.encode( | |
| pair, | |
| normalize_embeddings=True, | |
| convert_to_numpy=True, | |
| show_progress_bar=False, | |
| ) | |
| similarity_rows.append({ | |
| "Pair": pair_name, | |
| "Cosine similarity": round(float(np.dot(pair_vectors[0], pair_vectors[1])), 4), | |
| }) | |
| similarity_df = pd.DataFrame(similarity_rows) | |
| embedding_summary = pd.DataFrame([{ | |
| "Model": "all-MiniLM-L6-v2", | |
| "Device": DEVICE, | |
| "Chunks embedded": len(selected_chunks), | |
| "Vector dimension": int(vectors.shape[1]), | |
| "Output shape": str(tuple(vectors.shape)), | |
| "Embedding seconds": round(embedding_seconds, 3), | |
| "Chunks per second": round(len(selected_chunks) / max(embedding_seconds, 1e-9), 3), | |
| "Vector storage MB": round(vectors.nbytes / (1024 ** 2), 4), | |
| }]) | |
| figure, axes = plt.subplots(1, 2, figsize=(13, 5)) | |
| axes[0].hist(lengths, bins=min(15, max(5, len(lengths))), color="#4F46E5", edgecolor="white") | |
| axes[0].set_title("Chunk token distribution") | |
| axes[0].set_xlabel("Tokens per chunk") | |
| axes[0].set_ylabel("Number of chunks") | |
| axes[0].grid(alpha=0.2) | |
| axes[1].bar( | |
| similarity_df["Pair"], | |
| similarity_df["Cosine similarity"], | |
| color=["#16A34A", "#DC2626"], | |
| ) | |
| axes[1].set_title("Semantic similarity test") | |
| axes[1].set_xlabel("Sentence-pair type") | |
| axes[1].set_ylabel("Cosine similarity") | |
| axes[1].grid(axis="y", alpha=0.2) | |
| figure.tight_layout() | |
| document_summary.to_csv(run_directory / "document_summary.csv", index=False) | |
| chunk_summary.to_csv(run_directory / "chunk_summary.csv", index=False) | |
| pd.DataFrame({ | |
| "chunk_id": range(len(chunks)), | |
| "tokens": lengths, | |
| "text": chunks, | |
| }).to_csv(run_directory / "chunks.csv", index=False) | |
| embedding_summary.to_csv(run_directory / "embedding_summary.csv", index=False) | |
| similarity_df.to_csv(run_directory / "similarity_results.csv", index=False) | |
| np.save(run_directory / "vectors.npy", vectors) | |
| figure.savefig(run_directory / "benchmark_plots.png", dpi=180, bbox_inches="tight") | |
| settings = { | |
| "strategy": strategy, | |
| "chunk_size": int(chunk_size), | |
| "overlap_percent": effective_overlap, | |
| "model": MODEL_ID, | |
| "maximum_chunks": int(maximum_chunks), | |
| "device": DEVICE, | |
| "seed": SEED, | |
| } | |
| with open(run_directory / "settings.json", "w") as file: | |
| json.dump(settings, file, indent=2) | |
| archive_path = run_directory / "benchmark_results.zip" | |
| with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as archive: | |
| for output_path in run_directory.iterdir(): | |
| if output_path.is_file() and output_path != archive_path: | |
| archive.write(output_path, arcname=output_path.name) | |
| status = ( | |
| "### Benchmark completed\n" | |
| f"Created **{len(chunks)} chunks** and embedded " | |
| f"**{len(selected_chunks)} chunks** using **all-MiniLM-L6-v2**." | |
| ) | |
| return ( | |
| status, | |
| document_summary, | |
| chunk_summary, | |
| preview, | |
| embedding_summary, | |
| figure, | |
| str(archive_path), | |
| ) | |
| with gr.Blocks(title="Chunking and Embedding Benchmark") as demo: | |
| gr.Markdown( | |
| """ | |
| # Chunking and Embedding Benchmark | |
| Upload a document and test a chunking configuration. The application | |
| measures chunk quality, generates MiniLM embeddings and exports the results. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| uploaded_file = gr.File( | |
| label="Upload document", | |
| file_types=[".pdf", ".docx", ".txt", ".html", ".htm"], | |
| type="filepath", | |
| ) | |
| strategy = gr.Dropdown( | |
| choices=["Fixed", "Sentence", "Paragraph", "Recursive", "Semantic"], | |
| value="Recursive", | |
| label="Chunking strategy", | |
| ) | |
| chunk_size = gr.Dropdown( | |
| choices=[128, 256, 512, 1024], | |
| value=256, | |
| label="Target chunk size", | |
| ) | |
| overlap_percent = gr.Dropdown( | |
| choices=[0, 10, 20, 30], | |
| value=10, | |
| label="Overlap percentage", | |
| ) | |
| maximum_chunks = gr.Slider( | |
| minimum=5, | |
| maximum=300, | |
| value=100, | |
| step=5, | |
| label="Maximum chunks to embed", | |
| ) | |
| run_button = gr.Button("Run Benchmark", variant="primary") | |
| with gr.Column(scale=2): | |
| status = gr.Markdown("Upload a document and run the benchmark.") | |
| document_output = gr.Dataframe(label="Document summary", interactive=False) | |
| chunk_output = gr.Dataframe(label="Chunking results", interactive=False) | |
| with gr.Tab("Chunk preview"): | |
| preview_output = gr.Textbox(label="First three chunks", lines=22, interactive=False) | |
| with gr.Tab("Embedding results"): | |
| embedding_output = gr.Dataframe(label="Embedding measurements", interactive=False) | |
| plot_output = gr.Plot(label="Benchmark plots") | |
| with gr.Tab("Download"): | |
| download_output = gr.File(label="Download benchmark results") | |
| run_button.click( | |
| fn=run_benchmark, | |
| inputs=[uploaded_file, strategy, chunk_size, overlap_percent, maximum_chunks], | |
| outputs=[ | |
| status, | |
| document_output, | |
| chunk_output, | |
| preview_output, | |
| embedding_output, | |
| plot_output, | |
| download_output, | |
| ], | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch( | |
| theme=gr.themes.Soft(), | |
| ssr_mode=False, | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| ) |