From 8d83d5cfec2af9b3b1d61343921f6c6aed793fe7 Mon Sep 17 00:00:00 2001 From: Jonas Braus Date: Thu, 12 Dec 2024 15:35:32 +0100 Subject: [PATCH] chunker different methods --- chunker.py | 29 ++++++++++++++++++++++++++++- main.py | 3 ++- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/chunker.py b/chunker.py index fb4b7d3..61c2792 100644 --- a/chunker.py +++ b/chunker.py @@ -1,8 +1,11 @@ +import os + from llama_index.core import SimpleDirectoryReader from llama_index.core.node_parser import SemanticSplitterNodeParser from llama_index.embeddings.ollama import OllamaEmbedding +import json -def generate_chunks(content): +def generate_chunks_semantic(content): with open("./temp.txt", "wb") as file: file.write(content.encode("utf-8")) file.flush() @@ -21,3 +24,27 @@ def generate_chunks(content): output.append(node.to_dict()["text"]) return output + +def generate_chunks_line_split(content): + output = [] + + sp = content.split(".") + + while len(sp) > 0: + output.append("".join(sp[:10])) + sp = sp[10:] + + return output + + +def generate_chunks(content): + chunking_method = os.environ.get("CHUNKING_METHOD") + print("CHUNKING_METHOD:", chunking_method) + if chunking_method is None: + chunking_method = "none" + + return { + "semantic": lambda: generate_chunks_semantic(content), + "lines": lambda: generate_chunks_line_split(content), + "none": lambda: [content] + }[chunking_method.lower().strip()]() \ No newline at end of file diff --git a/main.py b/main.py index 291efbd..9c2c0ef 100644 --- a/main.py +++ b/main.py @@ -36,7 +36,8 @@ def retrieve_file_contents(path): contents = [] - for file in files: + print("parsing files...") + for file in tqdm(files): content = file_extraction_functions[file[0]](file[1]) with open(f"./{file[2]}.txt", "w", encoding="UTF-8") as future_s3_file: future_s3_file.writelines(file[1] + "\n" + content)