Spaces:
Sleeping
Sleeping
| """ | |
| embedder.py — Semantic Embedding Vector Generator & Visualizer. | |
| Features: | |
| - Generates per-word and per-sentence embeddings using sentence-transformers | |
| - PCA / t-SNE dimensionality reduction for 2D visualization | |
| - UMAP reduction for cluster visualization (if umap-learn installed) | |
| - Cosine semantic similarity heatmap between sentences | |
| - Word cluster labeling | |
| """ | |
| import numpy as np | |
| import streamlit as st | |
| from typing import List, Dict, Any, Optional, Tuple | |
| def get_embedding_model(model_name: str = "all-MiniLM-L6-v2"): | |
| """Load and cache the sentence-transformer embedding model.""" | |
| from sentence_transformers import SentenceTransformer | |
| return SentenceTransformer(model_name) | |
| def embed_texts( | |
| texts: List[str], | |
| model_name: str = "all-MiniLM-L6-v2", | |
| batch_size: int = 32 | |
| ) -> np.ndarray: | |
| """ | |
| Embed a list of texts into dense vectors. | |
| Returns ndarray of shape (N, embedding_dim). | |
| """ | |
| if not texts: | |
| return np.array([]) | |
| model = get_embedding_model(model_name) | |
| embeddings = model.encode( | |
| texts, | |
| batch_size=batch_size, | |
| show_progress_bar=False, | |
| normalize_embeddings=True, | |
| ) | |
| return embeddings | |
| def reduce_pca(embeddings: np.ndarray, n_components: int = 2) -> np.ndarray: | |
| """Reduce embedding dimensions using PCA.""" | |
| from sklearn.decomposition import PCA | |
| if len(embeddings) < n_components: | |
| return embeddings | |
| pca = PCA(n_components=n_components, random_state=42) | |
| return pca.fit_transform(embeddings) | |
| def reduce_tsne(embeddings: np.ndarray, n_components: int = 2, perplexity: int = None) -> np.ndarray: | |
| """Reduce embedding dimensions using t-SNE.""" | |
| from sklearn.manifold import TSNE | |
| n = len(embeddings) | |
| if n < 4: | |
| return reduce_pca(embeddings, n_components) | |
| perp = min(perplexity or 30, n - 1) | |
| tsne = TSNE( | |
| n_components=n_components, | |
| perplexity=perp, | |
| random_state=42, | |
| n_iter=500, | |
| learning_rate="auto", | |
| init="pca" | |
| ) | |
| return tsne.fit_transform(embeddings) | |
| def reduce_umap(embeddings: np.ndarray, n_components: int = 2) -> Optional[np.ndarray]: | |
| """Reduce embedding dimensions using UMAP (if available).""" | |
| try: | |
| import umap | |
| reducer = umap.UMAP(n_components=n_components, random_state=42) | |
| return reducer.fit_transform(embeddings) | |
| except ImportError: | |
| return None | |
| def cosine_similarity_matrix(embeddings: np.ndarray) -> np.ndarray: | |
| """Compute pairwise cosine similarity matrix (assumes normalized embeddings).""" | |
| norms = np.linalg.norm(embeddings, axis=1, keepdims=True) | |
| normalized = embeddings / np.clip(norms, 1e-8, None) | |
| return normalized @ normalized.T | |
| def cluster_words(embeddings: np.ndarray, n_clusters: int = 5) -> np.ndarray: | |
| """K-Means cluster the embeddings, return cluster labels.""" | |
| from sklearn.cluster import KMeans | |
| n = len(embeddings) | |
| k = min(n_clusters, max(1, n // 3)) | |
| kmeans = KMeans(n_clusters=k, random_state=42, n_init=10) | |
| return kmeans.fit_predict(embeddings) | |
| def build_embedding_data( | |
| ocr_results: List[Dict[str, Any]], | |
| text_unit: str = "word", # "word" | "sentence" | "line" | |
| reduction: str = "PCA", # "PCA" | "t-SNE" | "UMAP" | |
| n_clusters: int = 5, | |
| ) -> Dict[str, Any]: | |
| """ | |
| Full embedding pipeline: | |
| 1. Extract text units from OCR results | |
| 2. Generate embeddings | |
| 3. Reduce to 2D | |
| 4. Cluster | |
| 5. Return structured data for plotting | |
| Returns dict with keys: | |
| texts, embeddings, coords_2d, cluster_labels, similarity_matrix | |
| """ | |
| if not ocr_results: | |
| return {} | |
| # ── Step 1: Extract text units ── | |
| if text_unit == "word": | |
| texts = [r["text"] for r in ocr_results if r["text"].strip()] | |
| elif text_unit == "sentence": | |
| full = " ".join(r["text"] for r in ocr_results) | |
| # Simple sentence split | |
| import re | |
| sentences = re.split(r"(?<=[.!?])\s+", full) | |
| texts = [s.strip() for s in sentences if len(s.strip()) > 3] | |
| else: # line | |
| # Group words into lines by y-proximity | |
| sorted_results = sorted(ocr_results, key=lambda r: r["y_min"]) | |
| lines = [] | |
| current_line_y = None | |
| current_line_texts = [] | |
| for r in sorted_results: | |
| if current_line_y is None or abs(r["y_min"] - current_line_y) < 20: | |
| current_line_texts.append(r["text"]) | |
| current_line_y = r["y_min"] | |
| else: | |
| if current_line_texts: | |
| lines.append(" ".join(current_line_texts)) | |
| current_line_texts = [r["text"]] | |
| current_line_y = r["y_min"] | |
| if current_line_texts: | |
| lines.append(" ".join(current_line_texts)) | |
| texts = [l for l in lines if l.strip()] | |
| if not texts: | |
| return {} | |
| # ── Step 2: Embed ── | |
| embeddings = embed_texts(texts) | |
| if embeddings.size == 0: | |
| return {} | |
| # ── Step 3: Reduce to 2D ── | |
| if reduction == "t-SNE" and len(texts) >= 4: | |
| coords_2d = reduce_tsne(embeddings) | |
| elif reduction == "UMAP": | |
| coords_2d = reduce_umap(embeddings) | |
| if coords_2d is None: | |
| coords_2d = reduce_pca(embeddings) | |
| else: | |
| coords_2d = reduce_pca(embeddings) | |
| # ── Step 4: Cluster ── | |
| if len(texts) >= 2: | |
| cluster_labels = cluster_words(embeddings, n_clusters) | |
| else: | |
| cluster_labels = np.zeros(len(texts), dtype=int) | |
| # ── Step 5: Similarity matrix ── | |
| if len(texts) <= 60: | |
| sim_matrix = cosine_similarity_matrix(embeddings) | |
| else: | |
| sim_matrix = None | |
| # ── Step 6: Top-N similar pairs ── | |
| similar_pairs = [] | |
| if sim_matrix is not None and len(texts) > 1: | |
| for i in range(len(texts)): | |
| for j in range(i + 1, len(texts)): | |
| similar_pairs.append({ | |
| "i": i, | |
| "j": j, | |
| "text_i": texts[i], | |
| "text_j": texts[j], | |
| "similarity": float(sim_matrix[i, j]), | |
| }) | |
| similar_pairs.sort(key=lambda x: x["similarity"], reverse=True) | |
| similar_pairs = similar_pairs[:10] | |
| return { | |
| "texts": texts, | |
| "embeddings": embeddings, | |
| "coords_2d": coords_2d, | |
| "cluster_labels": cluster_labels.tolist(), | |
| "similarity_matrix": sim_matrix, | |
| "similar_pairs": similar_pairs, | |
| "embedding_dim": embeddings.shape[1] if len(embeddings.shape) > 1 else 0, | |
| "n_texts": len(texts), | |
| "reduction_method": reduction, | |
| } | |