diff --git a/chunker.py b/chunker.py index 4ad9c4f..fb4b7d3 100644 --- a/chunker.py +++ b/chunker.py @@ -9,7 +9,7 @@ def generate_chunks(content): file.close() embed_model = OllamaEmbedding(model_name="mxbai-embed-large") - splitter = SemanticSplitterNodeParser(buffer_size=5, breakpoint_percentile_threshold=70, embed_model=embed_model) + splitter = SemanticSplitterNodeParser(buffer_size=5, breakpoint_percentile_threshold=45, embed_model=embed_model) document = SimpleDirectoryReader(input_files=["./temp.txt"]).load_data() diff --git a/lib.py b/lib.py index 245f01d..90d3318 100644 --- a/lib.py +++ b/lib.py @@ -59,11 +59,20 @@ def prompt_embedding(prompt): ) cur = conn.cursor() - cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 3;", (embedding["embedding"],)) + cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 20;", (embedding["embedding"],)) result = cur.fetchall() - res = get_data_from_mongo(result) + initial_distance = result[0][2] + new_result = [] + + for res in result: + current_distance = res[2] + if abs(current_distance - initial_distance) < 1: + new_result.append(res) + + + res = get_data_from_mongo(new_result) print(ollama_conn.ask_ollama(prompt, res))