52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
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_semantic(content):
|
|
with open("./temp.txt", "wb") as file:
|
|
file.write(content.encode("utf-8"))
|
|
file.flush()
|
|
file.close()
|
|
|
|
embed_model = OllamaEmbedding(model_name="mxbai-embed-large")
|
|
splitter = SemanticSplitterNodeParser(buffer_size=5, breakpoint_percentile_threshold=45, embed_model=embed_model)
|
|
|
|
document = SimpleDirectoryReader(input_files=["./temp.txt"]).load_data()
|
|
|
|
nodes = splitter.build_semantic_nodes_from_documents(document)
|
|
|
|
os.remove("./temp.txt")
|
|
|
|
output = []
|
|
|
|
for i, node in enumerate(nodes):
|
|
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()]() |