itsprasun/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        chunks = []55        for i in range(0, len(text), self.chunk_size - self.chunk_overlap):56            chunks.append(text[i : i + self.chunk_size])57        return chunks58 59    def split_texts(self, texts: List[str]) -> List[str]:60        chunks = []61        for text in texts:62            chunks.extend(self.split(text))63        return chunks64 65 66class PDFLoader:67    def __init__(self, path: str):68        self.documents = []69        self.path = path70        print(f"PDFLoader initialized with path: {self.path}")71 72    def load(self):73        print(f"Loading PDF from path: {self.path}")74        print(f"Path exists: {os.path.exists(self.path)}")75        print(f"Is file: {os.path.isfile(self.path)}")76        print(f"Is directory: {os.path.isdir(self.path)}")77        print(f"File permissions: {oct(os.stat(self.path).st_mode)[-3:]}")78        79        try:80            # Try to open the file first to verify access81            with open(self.path, 'rb') as test_file:82                pass83            84            # If we can open it, proceed with loading85            self.load_file()86            87        except IOError as e:88            raise ValueError(f"Cannot access file at '{self.path}': {str(e)}")89        except Exception as e:90            raise ValueError(f"Error processing file at '{self.path}': {str(e)}")91 92    def load_file(self):93        with open(self.path, 'rb') as file:94            # Create PDF reader object95            pdf_reader = PyPDF2.PdfReader(file)96            97            # Extract text from each page98            text = ""99            for page in pdf_reader.pages:100                text += page.extract_text() + "\n"101            102            self.documents.append(text)103 104    def load_directory(self):105        for root, _, files in os.walk(self.path):106            for file in files:107                if file.lower().endswith('.pdf'):108                    file_path = os.path.join(root, file)109                    with open(file_path, 'rb') as f:110                        pdf_reader = PyPDF2.PdfReader(f)111                        112                        # Extract text from each page113                        text = ""114                        for page in pdf_reader.pages:115                            text += page.extract_text() + "\n"116                        117                        self.documents.append(text)118 119    def load_documents(self):120        self.load()121        return self.documents122 123 124if __name__ == "__main__":125    loader = TextFileLoader("data/KingLear.txt")126    loader.load()127    splitter = CharacterTextSplitter()128    chunks = splitter.split_texts(loader.documents)129    print(len(chunks))130    print(chunks[0])131    print("--------")132    print(chunks[1])133    print("--------")134    print(chunks[-2])135    print("--------")136    print(chunks[-1])137