Spaces:
Sleeping
Sleeping
| import os | |
| from typing import List | |
| class TextFileLoader: | |
| """ | |
| This class has the functionality to load the data from | |
| the text files. | |
| """ | |
| def __init__(self, path:str, encoding:str = "utf-8")->None: | |
| self.documents = [] | |
| self.path = path | |
| self.encoding = encoding | |
| def load(self)->None: | |
| """ | |
| if the path is of a directory, then load directory and read the file, | |
| else if the path is of the file, directly read the file. | |
| """ | |
| if os.path.isdir(self.path): | |
| self.load_directory() | |
| elif os.path.isfile(self.path) and self.path.endswith(".txt"): | |
| self.load_file() | |
| else: | |
| raise ValueError( | |
| "Provided path is neither a valid directory not a .txt tile" | |
| ) | |
| def load_file(self)->None: | |
| """ | |
| read the text file and append it to the list | |
| """ | |
| with open(self.path,"r",encoding=self.encoding) as f: | |
| self.documents.append(f.read()) | |
| def load_directory(self)->None: | |
| """ | |
| reads all the text files in the directory and appends it to the list | |
| """ | |
| for root, _, files in os.walk(self.path): | |
| for file in files: | |
| if file.endswith(".txt"): | |
| with open( | |
| os.path.join(root, file),"r",encoding=self.encoding | |
| ) as f: | |
| self.documents.append(f.read()) | |
| def load_documents(self): | |
| """ | |
| call the load function, that calls the function to read data and returns the documents. | |
| """ | |
| self.load() | |
| return self.documents | |
| class CharacterTextSplitter: | |
| """ | |
| This class contains the functionailites to chunk the text documents. | |
| """ | |
| def __init__(self, chunk_size:int = 1000,chunk_overlap:int = 200): | |
| assert(chunk_size>chunk_overlap),"Chunk size must be greater than chunk overlap" | |
| self.chunk_size = chunk_size | |
| self.chunk_overlap = chunk_overlap | |
| def split(self, text:str)->List[str]: | |
| """ | |
| takes in text and splits them based on character count | |
| """ | |
| chunks = [] | |
| for i in range(0, len(text),self.chunk_size-self.chunk_overlap): | |
| chunks.append(text[i:i+self.chunk_size]) | |
| return chunks | |
| def split_texts(self, texts:List[str])->List[str]: | |
| """ | |
| takes in list of texts and breaks it down to chunks | |
| """ | |
| chunks = [] | |
| for text in texts: | |
| chunks.extend(self.split(text)) | |
| return chunks | |
| if __name__ == "__main__": | |
| loader = TextFileLoader("/Users/shubham.agnihotri/Documents/GitHub/LLM-Ops/RAQA from scratch/data/KingLear.txt") | |
| loader.load() | |
| splitter = CharacterTextSplitter() | |
| chunks = splitter.split_texts(loader.documents) | |
| print(len(chunks)) | |
| print(chunks[0]) | |
| print("--------") | |
| print(chunks[1]) | |
| print("--------") | |
| print(chunks[-2]) | |
| print("--------") | |
| print(chunks[-1]) | |