From fde7e7570a3cd12c43edca430fb0cc1a4e6ac4ee Mon Sep 17 00:00:00 2001 From: Jonas Braus Date: Wed, 11 Dec 2024 22:21:14 +0100 Subject: [PATCH] puuuush --- lib.py | 2 +- main.py | 24 +++++++++++++----------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/lib.py b/lib.py index 0b2f1b7..1dc6247 100644 --- a/lib.py +++ b/lib.py @@ -55,7 +55,7 @@ def prompt_embedding(prompt): ) cur = conn.cursor() - cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 10;", (embedding["embedding"],)) + cur.execute(f"select id, doc_id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 3;", (embedding["embedding"],)) result = cur.fetchall() diff --git a/main.py b/main.py index c4c00df..95f31b0 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,8 @@ import ollama import pymupdf import os + +import chunker import lib import json import psycopg2 @@ -60,20 +62,20 @@ def create_embeddings(pContent): conn.commit() for content in pContent: - merged_info = "Dateiname: " + content["filename"] + " Dateiinhalt: " + content[ - "content"] # Bessere Embeddings mit Dateiname // Information vorne dran? - # print(merged_info) - response = ollama.embeddings(model="mxbai-embed-large", prompt=merged_info) - #embedding_list.append(response["embedding"]) - insert_data = f"insert into dbtable (filepath, embedding) Values ('{content['path']}', %s) Returning doc_id;" - cur.execute(insert_data, (response["embedding"],)) - doc_id = cur.fetchone()[0] - insert_data_mongo(doc_id, "Hello World") + for chunk in chunker.generate_chunks(content["content"]): + merged_info = "Dateiname: " + content["filename"] + " Dateiinhalt: " + chunk + # print(merged_info) + response = ollama.embeddings(model="mxbai-embed-large", prompt=merged_info) + #embedding_list.append(response["embedding"]) + insert_data = f"insert into dbtable (filepath, embedding) Values ('{content['path']}', %s) Returning doc_id;" + cur.execute(insert_data, (response["embedding"],)) + doc_id = cur.fetchone()[0] + insert_data_mongo(doc_id, chunk) conn.commit() conn.close() def insert_data_mongo(id, pChunk): - client = pymongo.MongoClient('mongodb://localhost:27017/') + client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/') mongodb = client['document_table'] collection = mongodb['documents'] dokument = { @@ -84,7 +86,7 @@ def insert_data_mongo(id, pChunk): def reload_files(): # path = input("Please provider path to folder: ") - contents = retrieve_file_contents("C:\\Users\\afist\\Downloads\\SoftwareEng") + contents = retrieve_file_contents("C:\\SoftwareEng") create_embeddings(contents)