Spaces:
Sleeping
Sleeping
| # from pypdf import PdfReader | |
| # import docx | |
| # from transformers.pipelines import pipeline | |
| # import streamlit as st | |
| # def extract_text(file): | |
| # text = "" | |
| # if file.name.endswith(".pdf"): | |
| # try: | |
| # reader = PdfReader(file) | |
| # for page in reader.pages: | |
| # text += page.extract_text() + "\n" | |
| # except Exception as e: | |
| # st.error(f"Error reading PDF {file.name}: {e}") | |
| # return "" | |
| # elif file.name.endswith(".docx"): | |
| # try: | |
| # document = docx.Document(file) | |
| # for paragraph in document.paragraphs: | |
| # text += paragraph.text + "\n" | |
| # except Exception as e: | |
| # st.error(f"Error reading DOCX {file.name}: {e}") | |
| # return "" | |
| # return text | |
| # def chunk_text(text, chunk_size=500, overlap=50): | |
| # chunks = [] | |
| # start = 0 | |
| # while start < len(text): | |
| # end = start + chunk_size | |
| # chunk = text[start:end] | |
| # chunks.append(chunk) | |
| # start = end - overlap | |
| # return chunks | |
| # def get_embeddings(texts): | |
| # try: | |
| # embedding_model = pipeline( | |
| # 'document-question-answering', | |
| # "sentence-transformers/all-MiniLM-L6-v2" | |
| # ) # Example model | |
| # embeddings = embedding_model(texts) | |
| # return embeddings | |
| # except Exception as e: | |
| # st.error(f"Error generating embeddings: {e}") | |
| # return [] | |
| # def process_files(files): | |
| # all_chunks = [] | |
| # all_embeddings = [] | |
| # chunks_metadata = [] | |
| # for file in files: | |
| # text = extract_text(file) | |
| # if not text: # Skip files that failed to process | |
| # continue | |
| # chunks = chunk_text(text) | |
| # embeddings = get_embeddings(chunks) | |
| # if not embeddings: # Skip files that failed to embed | |
| # continue | |
| # all_chunks.extend(chunks) | |
| # all_embeddings.extend(embeddings) | |
| # for i, chunk in enumerate(chunks): | |
| # chunks_metadata.append({"file_name": file.name, "chunk_index": i}) | |
| # print(f"Processed {len(files)} files, {len(all_chunks)} chunks generated.") | |
| # return all_chunks, all_embeddings, chunks_metadata | |
| import pypdf | |
| from docx import Document | |
| from transformers.pipelines import pipeline | |
| from sentence_transformers import SentenceTransformer | |
| import streamlit as st | |
| import numpy as np | |
| import os | |
| def extract_text(file): | |
| text = "" | |
| # Check if the input is a file path (string) or a file-like object | |
| if isinstance(file, str): | |
| file_name = os.path.basename(file) | |
| try: | |
| with open(file, 'rb') as f: # Open in binary mode | |
| if file_name.endswith(".pdf"): | |
| print('Processing pdf file.................\n') | |
| reader = pypdf.PdfReader(f) | |
| for page in reader.pages: | |
| text += page.extract_text() + "\\n" | |
| elif file_name.endswith(".docx"): | |
| document = Document(f) | |
| print('Processing DOCX file.................\n') | |
| for paragraph in document.paragraphs: | |
| if paragraph.text.strip(): # Check if the paragraph is not empty | |
| text += paragraph.text + "\\n" | |
| except FileNotFoundError: | |
| st.error(f"Error: File not found at {file}") | |
| return "" | |
| except Exception as e: | |
| st.error(f"Error reading {file_name}: {e}") | |
| return "" | |
| else: # Assume it's a file-like object (e.g., from Streamlit file_uploader) | |
| file_name = file.name | |
| try: | |
| if file_name.endswith(".pdf"): | |
| reader = pypdf.PdfReader(file) | |
| for page in reader.pages: | |
| text += page.extract_text() + "\\n" | |
| elif file_name.endswith(".docx"): | |
| document = Document(file) | |
| for paragraph in document.paragraphs: | |
| text += paragraph.text + "\\n" | |
| except Exception as e: | |
| st.error(f"Error reading {file_name}: {e}") | |
| return "" | |
| return text | |
| def chunk_text(text, chunk_size=1500, overlap=150): | |
| chunks = [] | |
| start = 0 | |
| while start < len(text): | |
| end = start + chunk_size | |
| chunk = text[start:end] | |
| chunks.append(chunk) | |
| start = end - overlap | |
| return chunks | |
| def get_embeddings(texts)-> np.ndarray: | |
| try: | |
| # embedding_model = pipeline( | |
| # "sentence-transformers/all-MiniLM-L6-v2" | |
| # ) # Example model | |
| # embeddings = embedding_model(texts) | |
| model = SentenceTransformer('dangvantuan/sentence-camembert-large') | |
| embeddings = model.encode(texts) | |
| print(f"Generated {len(embeddings)} embeddings.") | |
| return embeddings | |
| except Exception as e: | |
| st.error(f"Error generating embeddings: {e}") | |
| return [] | |
| def process_files(files): | |
| all_chunks = [] | |
| all_embeddings = [] | |
| chunks_metadata = [] | |
| for file in files: | |
| print(f"Processing file: {file.name if hasattr(file, 'name') else os.path.basename(file)}") | |
| text = extract_text(file) | |
| if not text: # Skip files that failed to process | |
| print(f"Skipping file {file.name if hasattr(file, 'name') else os.path.basename(file)} due to extraction error.") | |
| continue | |
| print(f"Chunking text...{file.name if hasattr(file, 'name') else os.path.basename(file)}\n") | |
| chunks = chunk_text(text) | |
| embeddings = get_embeddings(chunks) | |
| if not embeddings.any(): # Skip files that failed to embed | |
| continue | |
| all_chunks.extend(chunks) | |
| all_embeddings.extend(embeddings) | |
| for i, chunk in enumerate(chunks): | |
| chunks_metadata.append({"file_name": file.name if hasattr(file, 'name') else os.path.basename(file), "chunk_index": i}) | |
| return all_chunks, all_embeddings, chunks_metadata | |
| if __name__ == "__main__": | |
| # Example usage | |
| dummy_files = ['/Users/zac/Downloads/Janna/verbatimprocs/FZ- revenante - sept24.docx'] | |
| all_chunks, all_embeddings, chunks_metadata = process_files(dummy_files) | |
| print("Chunks ex:") | |
| print(f"Chunk 0: {all_chunks[0]}") | |
| print("\nEmbeddings:") | |
| print(f"Embedding 0: {all_embeddings[0][:10]}... (shape: {all_embeddings[0].shape})") # Print only the first 5 elements for brevity | |
| print("\nMetadata:") | |
| for i, metadata in enumerate(chunks_metadata[0:5]): | |
| print(f"Metadata {i}: {metadata}") | |
| """ | |
| Key improvements and explanations: | |
| Clear Function Definitions: Each function has a specific purpose with comprehensive docstrings. | |
| Error Handling: The extract_text and get_embeddings functions include try...except blocks to handle potential errors during file processing and embedding generation. Errors are displayed using st.error. | |
| File Type Handling: The extract_text function correctly handles both .pdf and .docx files. | |
| Chunking Strategy: The chunk_text function splits the text into smaller, overlapping chunks, which is a common strategy for RAG. | |
| Embedding Generation: The get_embeddings function uses the Hugging Face pipeline to generate embeddings. You can easily swap out the model if needed. | |
| Metadata: The process_files function now generates a list of metadata dictionaries, containing the file name and chunk index for each chunk. This is crucial for providing source attribution when answering queries. | |
| Testing: The if __name__ == "__main__": block provides example usage and testing of the functions. This is good practice for ensuring your code works as expected. I've added dummy file creation for testing. | |
| Efficiency: The code avoids unnecessary computations and handles files and text efficiently. | |
| """ | |