Spaces:
Paused
Paused
| import os | |
| from typing import List | |
| class TextFileLoader: | |
| def __init__(self, path: str, encoding: str = "utf-8"): | |
| self.documents = [] | |
| self.path = path | |
| self.encoding = encoding | |
| self.metadata = [] #added by YL | |
| def load(self): | |
| 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 nor a .txt file." | |
| ) | |
| def load_file(self): | |
| with open(self.path, "r", encoding=self.encoding) as f: | |
| self.documents.append(f.read()) | |
| self.metadata.append({"filename": os.path.basename(self.path), "filepath": self.path}) #added by YL | |
| def load_directory(self): | |
| 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_with_metadata(self): #changed name | |
| self.load() | |
| #return self.documents | |
| return [{"text": doc, "metadata": meta} for doc, meta in zip(self.documents, self.metadata)] #changed by YL | |
| class CharacterTextSplitter: | |
| 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]: | |
| # 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]: | |
| # chunks = [] | |
| # for text in texts: | |
| # chunks.extend(self.split(text)) | |
| # return chunks | |
| def split(self, text: str, metadata: dict[str, str]) -> List[dict[str, str]]: | |
| chunks = [] | |
| for i in range(0, len(text), self.chunk_size - self.chunk_overlap): | |
| chunk = text[i : i + self.chunk_size] | |
| chunk_metadata = metadata.copy() | |
| chunk_metadata["chunk_index"] = i // (self.chunk_size - self.chunk_overlap) | |
| chunks.append({"text": chunk, "metadata": chunk_metadata}) | |
| return chunks | |
| def split_texts_with_metadata( | |
| self, documents_with_metadata: List[dict[str, str]] | |
| ) -> List[dict[str, str]]: | |
| chunks = [] | |
| for document in documents_with_metadata: | |
| chunks.extend(self.split(document["text"], document["metadata"])) | |
| return chunks | |
| if __name__ == "__main__": | |
| loader = TextFileLoader("data/KingLear.txt") | |
| loader.load() | |
| splitter = CharacterTextSplitter() | |
| #chunks = splitter.split_texts(loader.documents) | |
| chunks = splitter.split_texts_with_metadata(loader.documents) | |
| print(len(chunks)) | |
| print(chunks[0]) | |
| print("--------") | |
| print(chunks[1]) | |
| print("--------") | |
| print(chunks[-2]) | |
| print("--------") | |
| print(chunks[-1]) |