atolat30/pythonic-rag-fastapi-react

1import os2from typing import List3import PyPDF24 5 6class TextFileLoader:7    def __init__(self, path: str, encoding: str = "utf-8"):8        self.documents = []9        self.path = path10        self.encoding = encoding11 12    def load(self):13        if os.path.isdir(self.path):14            self.load_directory()15        elif os.path.isfile(self.path) and self.path.endswith(".txt"):16            self.load_file()17        else:18            raise ValueError(19                "Provided path is neither a valid directory nor a .txt file."20            )21 22    def load_file(self):23        with open(self.path, "r", encoding=self.encoding) as f:24            self.documents.append(f.read())25 26    def load_directory(self):27        for root, _, files in os.walk(self.path):28            for file in files:29                if file.endswith(".txt"):30                    with open(31                        os.path.join(root, file), "r", encoding=self.encoding32                    ) as f:33                        self.documents.append(f.read())34 35    def load_documents(self):36        self.load()37        return self.documents38 39 40class CharacterTextSplitter:41    def __init__(42        self,43        chunk_size: int = 1000,44        chunk_overlap: int = 200,45    ):46        assert (47            chunk_size > chunk_overlap48        ), "Chunk size must be greater than chunk overlap"49 50        self.chunk_size = chunk_size51        self.chunk_overlap = chunk_overlap52 53    def split(self, text: str) -> List[str]:54        paragraphs = text.split('\n\n')55        chunks = []56        current_chunk = ""57        58        for paragraph in paragraphs:59            if len(current_chunk) + len(paragraph) > self.chunk_size:60                if current_chunk:61                    chunks.append(current_chunk.strip())62                if len(paragraph) > self.chunk_size:63                    words = paragraph.split()64                    current_chunk = ""65                    for word in words:66                        if len(current_chunk) + len(word) + 1 > self.chunk_size:67                            chunks.append(current_chunk.strip())68                            current_chunk = word69                        else:70                            current_chunk += " " + word if current_chunk else word71                else:72                    current_chunk = paragraph73            else:74                if current_chunk:75                    current_chunk += "\n\n" + paragraph76                else:77                    current_chunk = paragraph78        79        if current_chunk:80            chunks.append(current_chunk.strip())81            82        final_chunks = []83        for chunk in chunks:84            if len(chunk) > 8000:85                words = chunk.split()86                current = ""87                for word in words:88                    if len(current) + len(word) + 1 > 8000:89                        final_chunks.append(current.strip())90                        current = word91                    else:92                        current += " " + word if current else word93                if current:94                    final_chunks.append(current.strip())95            else:96                final_chunks.append(chunk)97                98        return final_chunks99 100    def split_texts(self, texts: List[str]) -> List[str]:101        chunks = []102        for text in texts:103            chunks.extend(self.split(text))104        return chunks105 106 107class PDFLoader:108    def __init__(self, path: str):109        self.documents = []110        self.path = path111        print(f"PDFLoader initialized with path: {self.path}")112 113    def load(self):114        print(f"Loading PDF from path: {self.path}")115        print(f"Path exists: {os.path.exists(self.path)}")116        print(f"Is file: {os.path.isfile(self.path)}")117        print(f"Is directory: {os.path.isdir(self.path)}")118        print(f"File permissions: {oct(os.stat(self.path).st_mode)[-3:]}")119        120        try:121            # Try to open the file first to verify access122            with open(self.path, 'rb') as test_file:123                pass124            125            # If we can open it, proceed with loading126            self.load_file()127            128        except IOError as e:129            print(f"IOError while accessing file: {str(e)}")130            raise ValueError(f"Cannot access file at '{self.path}': {str(e)}")131        except Exception as e:132            print(f"Unexpected error while processing file: {str(e)}")133            raise ValueError(f"Error processing file at '{self.path}': {str(e)}")134 135    def load_file(self):136        with open(self.path, 'rb') as file:137            # Create PDF reader object138            pdf_reader = PyPDF2.PdfReader(file)139            print(f"PDF loaded successfully. Number of pages: {len(pdf_reader.pages)}")140            141            # Extract text from each page142            text = ""143            for i, page in enumerate(pdf_reader.pages):144                page_text = page.extract_text()145                if not page_text.strip():146                    print(f"Warning: Page {i+1} appears to be empty or unreadable")147                text += page_text + "\n"148                print(f"Processed page {i+1}, extracted {len(page_text)} characters")149            150            if not text.strip():151                print("Warning: No text was extracted from the PDF")152            else:153                print(f"Successfully extracted {len(text)} characters of text")154            155            self.documents.append(text)156 157    def load_directory(self):158        for root, _, files in os.walk(self.path):159            for file in files:160                if file.lower().endswith('.pdf'):161                    file_path = os.path.join(root, file)162                    with open(file_path, 'rb') as f:163                        pdf_reader = PyPDF2.PdfReader(f)164                        165                        # Extract text from each page166                        text = ""167                        for page in pdf_reader.pages:168                            text += page.extract_text() + "\n"169                        170                        self.documents.append(text)171 172    def load_documents(self):173        self.load()174        return self.documents175 176 177if __name__ == "__main__":178    loader = TextFileLoader("data/KingLear.txt")179    loader.load()180    splitter = CharacterTextSplitter()181    chunks = splitter.split_texts(loader.documents)182    print(len(chunks))183    print(chunks[0])184    print("--------")185    print(chunks[1])186    print("--------")187    print(chunks[-2])188    print("--------")189    print(chunks[-1])190