Spaces:
Running
Running
| import PyPDF2 | |
| from docx import Document | |
| import markdown | |
| import re | |
| class DocumentProcessor: | |
| def __init__(self, chunk_size=1000, overlap=100): | |
| self.chunk_size = chunk_size | |
| self.overlap = overlap | |
| def extract_text(self, filepath, file_type): | |
| if file_type == 'txt': | |
| return self._extract_from_txt(filepath) | |
| elif file_type == 'pdf': | |
| return self._extract_from_pdf(filepath) | |
| elif file_type == 'docx': | |
| return self._extract_from_docx(filepath) | |
| elif file_type == 'md': | |
| return self._extract_from_markdown(filepath) | |
| else: | |
| raise ValueError(f"Unsupported file type: {file_type}") | |
| def _extract_from_txt(self, filepath): | |
| with open(filepath, 'r', encoding='utf-8') as file: | |
| return file.read() | |
| def _extract_from_pdf(self, filepath): | |
| text = "" | |
| with open(filepath, 'rb') as file: | |
| pdf_reader = PyPDF2.PdfReader(file) | |
| for page_num in range(len(pdf_reader.pages)): | |
| page = pdf_reader.pages[page_num] | |
| text += page.extract_text() + "\n" | |
| return text | |
| def _extract_from_docx(self, filepath): | |
| doc = Document(filepath) | |
| text = "" | |
| for paragraph in doc.paragraphs: | |
| text += paragraph.text + "\n" | |
| return text | |
| def _extract_from_markdown(self, filepath): | |
| with open(filepath, 'r', encoding='utf-8') as file: | |
| md_text = file.read() | |
| html = markdown.markdown(md_text) | |
| text = re.sub('<[^<]+?>', '', html) | |
| return text | |
| def split_into_chunks(self, text): | |
| sentences = self._split_into_sentences(text) | |
| chunks = [] | |
| current_chunk = "" | |
| for sentence in sentences: | |
| if len(current_chunk) + len(sentence) <= self.chunk_size: | |
| current_chunk += sentence + " " | |
| else: | |
| if current_chunk: | |
| chunks.append(current_chunk.strip()) | |
| current_chunk = sentence + " " | |
| if current_chunk: | |
| chunks.append(current_chunk.strip()) | |
| overlapped_chunks = [] | |
| for i, chunk in enumerate(chunks): | |
| if i > 0 and self.overlap > 0: | |
| prev_sentences = chunks[i-1].split('.')[-2:] | |
| overlap_text = '. '.join(prev_sentences).strip() | |
| if overlap_text: | |
| chunk = overlap_text + ". " + chunk | |
| overlapped_chunks.append(chunk) | |
| return overlapped_chunks | |
| def _split_into_sentences(self, text): | |
| sentence_endings = re.compile(r'[.!?]+') | |
| sentences = sentence_endings.split(text) | |
| sentences = [s.strip() for s in sentences if s.strip()] | |
| result = [] | |
| for i, sentence in enumerate(sentences): | |
| if i < len(sentences) - 1: | |
| next_char_match = sentence_endings.search(text, text.find(sentence) + len(sentence)) | |
| if next_char_match: | |
| sentence += next_char_match.group() | |
| result.append(sentence) | |
| return result |