Update doc_preprocessing.py

doc_preprocessing.py

@@ -1,70 +1,3 @@
-# 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
@@ -140,26 +73,50 @@ def get_embeddings(texts)-> np.ndarray:
         st.error(f"Error generating embeddings: {e}")
         return []
 
-def process_files(files):
+def process_files(uploaded_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: # Skip files that failed to embed
-        #     continue
+    for uploaded_file in uploaded_files:
+        # Create a temporary file within a writable directory
+        # Hugging Face Spaces usually allows writing to /tmp/ or your app's directory
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{uploaded_file.type.split('/')[-1]}") as temp_file:
+            temp_file.write(uploaded_file.getvalue())
+            temp_file_path = temp_file.name
 
-        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
+        try:
+            # Now, use temp_file_path to process the file
+            # Your existing process_files logic would go here,
+            # reading from temp_file_path
+            st.write(f"Processing file: {temp_file_path}")
+            # Example: Replace this with your actual processing logic
+            # For PDF/Word, you'd likely use a library like pypdf, python-docx, or langchain loaders
+            if uploaded_file.type == "application/pdf":
+                # Example for PDF:
+                # from pypdf import PdfReader
+                # reader = PdfReader(temp_file_path)
+                # text = ""
+                # for page in reader.pages:
+                #     text += page.extract_text() + "\n"
+                pass # Replace with actual PDF processing
+            elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+                # Example for DOCX:
+                # from docx import Document
+                # document = Document(temp_file_path)
+                # text = ""
+                # for paragraph in document.paragraphs:
+                #     text += paragraph.text + "\n"
+                pass # Replace with actual DOCX processing
+
+            # Placeholder for chunking, embedding, and metadata extraction
+            # (You'd replace this with your actual logic based on the processed file content)
+            all_chunks.append(f"Content from {uploaded_file.name}")
+            all_embeddings.append([0.1, 0.2, 0.3]) # Replace with actual embeddings
+            chunks_metadata.append({"filename": uploaded_file.name})
+
+        finally:
+            # Clean up the temporary file
+            os.remove(temp_file_path)
 
+    return all_chunks, all_embeddings, chunks_metadata