rosa0003/smartpdf_Highlighter

1"""2This module provides functions for generating a highlighted PDF with important sentences.3 4The main function, `generate_highlighted_pdf`, takes an input PDF file and a pre-trained5sentence embedding model as input.6 7It splits the text of the PDF into sentences, computes sentence embeddings, and builds a8graph based on the cosine similarity between embeddings and at the same time split the9sentences to different clusters using clustering.10 11The sentences are then ranked using PageRank scores and a the middle of the cluster,12and important sentences are selected based on a threshold and clustering.13 14Finally, the selected sentences are highlighted in the PDF and the highlighted PDF content15is returned.16 17Other utility functions in this module include functions for loading a sentence embedding18model, encoding sentences, computing similarity matrices,building graphs, ranking sentences,19clustering sentence embeddings, and splitting text into sentences.20 21Note: This module requires the PyMuPDF, networkx, numpy, torch, sentence_transformers, and22sklearn libraries to be installed.23"""24 25import logging26from typing import BinaryIO, List, Tuple27 28import fitz  # PyMuPDF29import networkx as nx30import numpy as np31import torch32import torch.nn.functional as F33from sentence_transformers import SentenceTransformer34from sklearn.cluster import KMeans35 36# Constants37MAX_PAGE = 4038MAX_SENTENCES = 200039PAGERANK_THRESHOLD_RATIO = 0.1540NUM_CLUSTERS_RATIO = 0.0541MIN_WORDS = 1042 43# Logger configuration44logging.basicConfig(level=logging.ERROR)45logger = logging.getLogger(__name__)46 47 48def load_sentence_model(revision: str = None) -> SentenceTransformer:49    """50    Load a pre-trained sentence embedding model.51 52    Args:53        revision (str): Optional parameter to specify the model revision.54 55    Returns:56        SentenceTransformer: A pre-trained sentence embedding model.57    """58    return SentenceTransformer("avsolatorio/GIST-Embedding-v0", revision=revision)59 60 61def encode_sentence(model: SentenceTransformer, sentence: str) -> torch.Tensor:62    """63    Encode a sentence into a fixed-dimensional vector representation.64 65    Args:66        model (SentenceTransformer): A pre-trained sentence embedding model.67        sentence (str): Input sentence.68 69    Returns:70        torch.Tensor: Encoded sentence vector.71    """72 73    model.eval()  # Set the model to evaluation mode74 75    # Check if GPU is available76    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")77 78    with torch.no_grad():  # Disable gradient tracking79        return model.encode(sentence, convert_to_tensor=True).to(device)80 81 82def compute_similarity_matrix(embeddings: torch.Tensor) -> np.ndarray:83    """84    Compute the cosine similarity matrix between sentence embeddings.85 86    Args:87        embeddings (torch.Tensor): Sentence embeddings.88 89    Returns:90        np.ndarray: Cosine similarity matrix.91    """92    scores = F.cosine_similarity(93        embeddings.unsqueeze(1), embeddings.unsqueeze(0), dim=-194    )95    similarity_matrix = scores.cpu().numpy()96    normalized_adjacency_matrix = similarity_matrix / similarity_matrix.sum(97        axis=1, keepdims=True98    )99    return normalized_adjacency_matrix100 101 102def build_graph(normalized_adjacency_matrix: np.ndarray) -> nx.DiGraph:103    """104    Build a directed graph from a normalized adjacency matrix.105 106    Args:107        normalized_adjacency_matrix (np.ndarray): Normalized adjacency matrix.108 109    Returns:110        nx.DiGraph: Directed graph.111    """112    return nx.DiGraph(normalized_adjacency_matrix)113 114 115def rank_sentences(graph: nx.DiGraph, sentences: List[str]) -> List[Tuple[str, float]]:116    """117    Rank sentences based on PageRank scores.118 119    Args:120        graph (nx.DiGraph): Directed graph.121        sentences (List[str]): List of sentences.122 123    Returns:124        List[Tuple[str, float]]: Ranked sentences with their PageRank scores.125    """126    pagerank_scores = nx.pagerank(graph)127    ranked_sentences = sorted(128        zip(sentences, pagerank_scores.values()),129        key=lambda x: x[1],130        reverse=True,131    )132    return ranked_sentences133 134 135def cluster_sentences(136    embeddings: torch.Tensor, num_clusters: int137) -> Tuple[np.ndarray, np.ndarray]:138    """139    Cluster sentence embeddings using K-means clustering.140 141    Args:142        embeddings (torch.Tensor): Sentence embeddings.143        num_clusters (int): Number of clusters.144 145    Returns:146        Tuple[np.ndarray, np.ndarray]: Cluster assignments and cluster centers.147    """148    kmeans = KMeans(n_clusters=num_clusters, random_state=42)149    cluster_assignments = kmeans.fit_predict(embeddings.cpu())150    cluster_centers = kmeans.cluster_centers_151    return cluster_assignments, cluster_centers152 153 154def get_middle_sentence(cluster_indices: np.ndarray, sentences: List[str]) -> List[str]:155    """156    Get the middle sentence from each cluster.157 158    Args:159        cluster_indices (np.ndarray): Cluster assignments.160        sentences (List[str]): List of sentences.161 162    Returns:163        List[str]: Middle sentences from each cluster.164    """165    middle_indices = [166        int(np.median(np.where(cluster_indices == i)[0]))167        for i in range(max(cluster_indices) + 1)168    ]169    middle_sentences = [sentences[i] for i in middle_indices]170    return middle_sentences171 172 173def split_text_into_sentences(text: str, min_words: int = MIN_WORDS) -> List[str]:174    """175    Split text into sentences.176 177    Args:178        text (str): Input text.179        min_words (int): Minimum number of words for a valid sentence.180 181    Returns:182        List[str]: List of sentences.183    """184    sentences = []185    for s in text.split("."):186        s = s.strip()187        # filtering out short sentences and sentences that contain more than 40% digits188        if (189            s190            and len(s.split()) >= min_words191            and (sum(c.isdigit() for c in s) / len(s)) < 0.4192        ):193            sentences.append(s)194    return sentences195 196 197def extract_text_from_pages(doc):198    """Generator to yield text per page from the PDF, for memory efficiency for large PDFs."""199    for page_num in range(len(doc)):200        yield doc[page_num].get_text()201 202 203def generate_highlighted_pdf(204    input_pdf_file: BinaryIO, model=load_sentence_model()205) -> bytes:206    """207    Generate a highlighted PDF with important sentences.208 209    Args:210        input_pdf_file: Input PDF file object.211        model (SentenceTransformer): Pre-trained sentence embedding model.212 213    Returns:214        bytes: Highlighted PDF content.215    """216    with fitz.open(stream=input_pdf_file.read(), filetype="pdf") as doc:217        num_pages = doc.page_count218 219        if num_pages > MAX_PAGE:220            # It will show the error message for the user.221            return f"The PDF file exceeds the maximum limit of {MAX_PAGE} pages."222 223        sentences = []224        for page_text in extract_text_from_pages(doc):  # Memory efficient225            sentences.extend(split_text_into_sentences(page_text))226 227        len_sentences = len(sentences)228 229        print(len_sentences)230 231        if len_sentences > MAX_SENTENCES:232            # It will show the error message for the user.233            return (234                f"The PDF file exceeds the maximum limit of {MAX_SENTENCES} sentences."235            )236 237        embeddings = encode_sentence(model, sentences)238        similarity_matrix = compute_similarity_matrix(embeddings)239        graph = build_graph(similarity_matrix)240        ranked_sentences = rank_sentences(graph, sentences)241 242        pagerank_threshold = int(len(ranked_sentences) * PAGERANK_THRESHOLD_RATIO) + 1243        top_pagerank_sentences = [244            sentence[0] for sentence in ranked_sentences[:pagerank_threshold]245        ]246 247        num_clusters = int(len_sentences * NUM_CLUSTERS_RATIO) + 1248        cluster_assignments, _ = cluster_sentences(embeddings, num_clusters)249 250        center_sentences = get_middle_sentence(cluster_assignments, sentences)251        important_sentences = list(set(top_pagerank_sentences + center_sentences))252 253        for i in range(num_pages):254            try:255                page = doc[i]256 257                for sentence in important_sentences:258                    rects = page.search_for(sentence)259                    colors = (fitz.pdfcolor["yellow"], fitz.pdfcolor["green"])260 261                    for i, rect in enumerate(rects):262                        color = colors[i % 2]263                        annot = page.add_highlight_annot(rect)264                        annot.set_colors(stroke=color)265                        annot.update()266            except Exception as e:267                logger.error(f"Error processing page {i}: {e}")268 269        output_pdf = doc.write()270 271    return output_pdf272