huggingchat/document-parser-rag

1import gradio as gr2import spaces3import subprocess4import os5import shutil6import string7import random8from pypdf import PdfReader9import ocrmypdf10from sentence_transformers import SentenceTransformer11 12model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m")13model.to(device="cuda")14 15 16@spaces.GPU17def embed(queries, chunks) -> dict[str, list[tuple[str, float]]]:18    query_embeddings = model.encode(queries, prompt_name="query")19    document_embeddings = model.encode(chunks)20 21    scores = query_embeddings @ document_embeddings.T22    results = {}23    for query, query_scores in zip(queries, scores):24        chunk_idxs = [i for i in range(len(chunks))]25        # Get a structure like {query: [(chunk_idx, score), (chunk_idx, score), ...]}26        results[query] = list(zip(chunk_idxs, query_scores))27 28    return results29 30 31def random_word(length):32    letters = string.ascii_lowercase33    return "".join(random.choice(letters) for _ in range(length))34 35 36def convert_pdf(input_file) -> str:37    reader = PdfReader(input_file)38    text = extract_text_from_pdf(reader)39 40    # Check if there are any images41    image_count = 042    for page in reader.pages:43        image_count += len(page.images)44 45    # If there are images and not much content, perform OCR on the document46    if image_count > 0 and len(text) < 1000:47        out_pdf_file = input_file.replace(".pdf", "_ocr.pdf")48        ocrmypdf.ocr(input_file, out_pdf_file, force_ocr=True)49 50        # Re-extract text51        text = extract_text_from_pdf(PdfReader(input_file))52 53        # Delete the OCR file54        os.remove(out_pdf_file)55 56    return text57 58 59def extract_text_from_pdf(reader):60    full_text = ""61    for idx, page in enumerate(reader.pages):62        text = page.extract_text()63        if len(text) > 0:64            full_text += f"---- Page {idx} ----\n" + page.extract_text() + "\n\n"65 66    return full_text.strip()67 68 69def convert_pandoc(input_file, filename) -> str:70    # Temporarily copy the file71    shutil.copyfile(input_file, filename)72 73    # Convert the file to markdown with pandoc74    output_file = f"{random_word(16)}.md"75    result = subprocess.call(["pandoc", filename, "-t", "markdown", "-o", output_file])76    if result != 0:77        raise ValueError("Error converting file to markdown with pandoc")78 79    # Read the file and delete temporary files80    with open(output_file, "r") as f:81        markdown = f.read()82    os.remove(output_file)83    os.remove(filename)84 85    return markdown86 87 88@spaces.GPU89def convert(input_file, filename) -> str:90    plain_text_filetypes = [91        ".txt",92        ".csv",93        ".tsv",94        ".md",95        ".yaml",96        ".toml",97        ".json",98        ".json5",99        ".jsonc",100    ]101    # Already a plain text file that wouldn't benefit from pandoc so return the content102    if any(filename.endswith(ft) for ft in plain_text_filetypes):103        with open(input_file, "r") as f:104            return f.read()105 106    if filename.endswith(".pdf"):107        return convert_pdf(input_file)108 109    return convert_pandoc(input_file, filename)110 111 112def chunk_to_length(text, max_length=512):113    chunks = []114    while len(text) > max_length:115        chunks.append(text[:max_length])116        text = text[max_length:]117    chunks.append(text)118    return chunks119 120 121@spaces.GPU122def predict(queries, documents, document_filenames, max_characters) -> list[list[str]]:123    queries = queries.split("\n")124    document_filenames = document_filenames.split("\n")125 126    # Convert the documents to text127    converted_docs = [128        convert(doc, filename) for doc, filename in zip(documents, document_filenames)129    ]130 131    # Return if the total length is less than the max characters132    total_doc_lengths = sum([len(doc) for doc in converted_docs])133    if total_doc_lengths < max_characters:134        return [[doc] for doc, _ in converted_docs]135 136    # Embed the documents in 512 character chunks137    chunked_docs = [chunk_to_length(doc, 512) for doc in converted_docs]138    embedded_docs = [embed(queries, chunks) for chunks in chunked_docs]139 140    # Get a structure like {query: [(doc_idx, chunk_idx, score), (doc_idx, chunk_idx, score), ...]}141    query_embeddings = {}142    for doc_idx, embedded_doc in enumerate(embedded_docs):143        for query, doc_scores in embedded_doc.items():144            doc_scores_with_doc = [145                (doc_idx, chunk_idx, score) for (chunk_idx, score) in doc_scores146            ]147            if query not in query_embeddings:148                query_embeddings[query] = []149            query_embeddings[query] = query_embeddings[query] + doc_scores_with_doc150 151    # Sort the embeddings by score152    for query, doc_scores in query_embeddings.items():153        query_embeddings[query] = sorted(doc_scores, key=lambda x: x[2], reverse=True)154 155    # Choose the top embedding from each query until we reach the max characters156    # Getting a structure like [[chunk, ...]]157    document_embeddings = [[] for _ in range(len(documents))]158    total_chars = 0159    while (160        total_chars < max_characters161        and sum([len(x) for x in query_embeddings.values()]) > 0162    ):163        for query, doc_scores in query_embeddings.items():164            if len(doc_scores) == 0:165                continue166 167            # Grab the top score for the query168            doc_idx, chunk_idx, _ = doc_scores.pop(0)169 170            # Ensure we have space171            chunk = chunked_docs[doc_idx][chunk_idx]172            if total_chars + len(chunk) > max_characters:173                continue174 175            # Ensure we haven't already added this chunk from this document176            if chunk_idx in document_embeddings[doc_idx]:177                continue178 179            # Add the chunk180            document_embeddings[doc_idx].append(chunk_idx)181            total_chars += len(chunk)182 183    # Get the actual text for the chunks184    document_embeddings = [185        [chunked_docs[doc_idx][chunk_idx] for chunk_idx in chunks]186        for doc_idx, chunks in enumerate(document_embeddings)187    ]188 189    return document_embeddings190 191 192# We accept a filename because the gradio JS interface removes this information193# and it's critical for choosing the correct processing pipeline194gr.Interface(195    predict,196    inputs=[197        gr.Textbox(label="Queries separated by newline"),198        gr.File(label="Upload File", file_count="multiple"),199        gr.Textbox(label="Filenames separated by newline"),200        gr.Number(label="Max output characters", value=16384),201    ],202    outputs=[gr.JSON(label="Embedded documents")],203).launch()204