Spaces:
Sleeping
Sleeping
| # Document processing utilities for different file types | |
| import os | |
| import tempfile | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| import fitz # PyMuPDF | |
| from docx import Document | |
| import streamlit as st | |
| from config import EMBEDDING_DIMENSIONS | |
| from device_setup import clear_gpu_memory | |
| class DocumentProcessor: | |
| """Handles processing of different document types""" | |
| def __init__(self, clip_model, clip_processor, whisper_model, ocr_reader, text_encoder, device): | |
| self.clip_model = clip_model | |
| self.clip_processor = clip_processor | |
| self.whisper_model = whisper_model | |
| self.ocr_reader = ocr_reader | |
| self.text_encoder = text_encoder | |
| self.device = device | |
| def extract_text_from_pdf(self, pdf_path: str) -> str: | |
| """extract text from PDF""" | |
| try: | |
| doc = fitz.open(pdf_path) | |
| text = "" | |
| for page in doc: | |
| text += page.get_text() | |
| doc.close() | |
| return text.strip() | |
| except Exception as e: | |
| st.error(f"Error extracting PDF text: {e}") | |
| return "" | |
| def extract_text_from_docx(self, docx_path: str) -> str: | |
| """extract text content from DOCX""" | |
| try: | |
| doc = Document(docx_path) | |
| text = "" | |
| for paragraph in doc.paragraphs: | |
| text += paragraph.text + "\n" | |
| return text.strip() | |
| except Exception as e: | |
| st.error(f"Error extracting DOCX text: {e}") | |
| return "" | |
| def extract_text_from_image(self, image_path: str) -> str: | |
| """extract text from image using OCR""" | |
| try: | |
| results = self.ocr_reader.readtext(image_path) | |
| text = " ".join([result[1] for result in results]) | |
| return text.strip() | |
| except Exception as e: | |
| st.error(f"Error extracting text from image: {e}") | |
| return "" | |
| def get_image_embedding(self, image_path: str) -> np.ndarray: | |
| """generate CLIP embedding for image""" | |
| try: | |
| image = Image.open(image_path).convert("RGB") | |
| inputs = self.clip_processor(images=image, return_tensors="pt", padding=True) | |
| inputs = {k: v.to(self.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| image_features = self.clip_model.get_image_features(**inputs) | |
| embedding = image_features.cpu().numpy().flatten() | |
| clear_gpu_memory() | |
| return embedding | |
| except Exception as e: | |
| st.error(f"Error getting image embedding: {e}") | |
| # Return zero embedding with same dimension | |
| clip_dim = EMBEDDING_DIMENSIONS.get("clip_base", 512) | |
| return np.zeros(clip_dim) | |
| def transcribe_audio(self, audio_path: str) -> str: | |
| """transcribe audio file to text using Whisper""" | |
| try: | |
| result = self.whisper_model.transcribe(audio_path) | |
| return result["text"].strip() | |
| except Exception as e: | |
| st.error(f"Error transcribing audio: {e}") | |
| return "" | |
| def get_text_embedding(self, text: str) -> np.ndarray: | |
| """generate text embedding using sentence transformer""" | |
| try: | |
| if not text.strip(): | |
| return np.zeros(EMBEDDING_DIMENSIONS["text"]) | |
| with torch.no_grad(): | |
| embedding = self.text_encoder.encode([text], convert_to_numpy=True) | |
| return embedding[0] | |
| except Exception as e: | |
| st.error(f"Error getting text embedding: {e}") | |
| return np.zeros(EMBEDDING_DIMENSIONS["text"]) | |
| def process_uploaded_file(self, uploaded_file, file_type: str, data_dir) -> dict: | |
| """ | |
| Process uploaded file and return document dictionary | |
| Returns None if processing fails | |
| """ | |
| try: | |
| # save uploaded file to temporary location | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_type}") as tmp_file: | |
| tmp_file.write(uploaded_file.read()) | |
| tmp_path = tmp_file.name | |
| # process based on file type | |
| if file_type == 'pdf': | |
| content = self.extract_text_from_pdf(tmp_path) | |
| embedding = self.get_text_embedding(content) | |
| embedding_type = 'text' | |
| elif file_type == 'docx': | |
| content = self.extract_text_from_docx(tmp_path) | |
| embedding = self.get_text_embedding(content) | |
| embedding_type = 'text' | |
| elif file_type in ['png', 'jpg', 'jpeg']: | |
| content = self.extract_text_from_image(tmp_path) | |
| embedding = self.get_image_embedding(tmp_path) | |
| embedding_type = 'image' | |
| elif file_type in ['mp3', 'wav', 'ogg', 'm4a']: | |
| content = self.transcribe_audio(tmp_path) | |
| embedding = self.get_text_embedding(content) | |
| embedding_type = 'text' | |
| else: | |
| os.unlink(tmp_path) | |
| return None | |
| # clean up temp file | |
| os.unlink(tmp_path) | |
| return { | |
| 'content': content, | |
| 'embedding': embedding, | |
| 'embedding_type': embedding_type | |
| } | |
| except Exception as e: | |
| st.error(f"Error processing file {uploaded_file.name}: {e}") | |
| if 'tmp_path' in locals() and os.path.exists(tmp_path): | |
| os.unlink(tmp_path) | |
| return None |