Spaces:
Sleeping
Sleeping
File size: 3,041 Bytes
5623f53 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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])
|