From cda757029e0dd5bd15d616adce58e8cb6bb66ce9 Mon Sep 17 00:00:00 2001 From: Jonas Braus Date: Tue, 10 Dec 2024 23:34:44 +0100 Subject: [PATCH] puuuuush --- lib.py | 4 +++- main.py | 39 +++++++++++++++++++++--------- main2.py | 60 ----------------------------------------------- test_PDFparser.py | 13 ---------- test_get_files.py | 24 ------------------- test_image.py | 7 ------ test_mp3parser.py | 12 ---------- 7 files changed, 31 insertions(+), 128 deletions(-) delete mode 100644 main2.py delete mode 100644 test_PDFparser.py delete mode 100644 test_get_files.py delete mode 100644 test_image.py delete mode 100644 test_mp3parser.py diff --git a/lib.py b/lib.py index 489d195..0b2f1b7 100644 --- a/lib.py +++ b/lib.py @@ -55,7 +55,9 @@ def prompt_embedding(prompt): ) cur = conn.cursor() - cur.execute(f"select id, filepath, embedding, embedding <-> {embedding} as distance from dbtable order by distance limit 1;") + cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 10;", (embedding["embedding"],)) result = cur.fetchall() + + conn.close() print(result) \ No newline at end of file diff --git a/main.py b/main.py index aa50c6d..b63d1fc 100644 --- a/main.py +++ b/main.py @@ -31,10 +31,6 @@ def retrieve_file_contents(path): return contents -contents = retrieve_file_contents("C:\\Users\\afist\\Downloads\\SoftwareEng") - - -embedding_list = [] def create_embeddings(pContent): conn = psycopg2.connect( @@ -49,35 +45,56 @@ def create_embeddings(pContent): create_table_query = ''' create table if not exists dbtable ( id SERIAL PRIMARY KEY, - filename TEXT NOT NULL, - embedding VECTOR NOT NULL - ) + filepath TEXT NOT NULL, + embedding VECTOR NOT NULL + ); ''' + cur.execute('CREATE EXTENSION IF NOT EXISTS vector;') cur.execute(create_table_query) - cur.execute('CREATE EXTENSION IF NOT EXISTS vector;') + 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["filepath"]}, {response["embedding"]});' - cur.execute(insert_data) + insert_data = f"insert into dbtable (filepath, embedding) Values ('{content['path']}', %s);" + cur.execute(insert_data, (response["embedding"],)) conn.commit() conn.close() +def reload_files(): + # path = input("Please provider path to folder: ") + contents = retrieve_file_contents("C:\\SoftwareEng") + create_embeddings(contents) +def add_files(): + path = input("Please provider path to folder: ") + pass +def prompt_cycle(): + while True: + prompt = input("Please enter prompt: ") + lib.prompt_embedding(prompt) -create_embeddings(contents) +def get_user_action(): + user_action = input("Reload Files (r), Add Files (a), Prompt (p): ") + { + "r": lambda: reload_files(), + "a": lambda: add_files(), + "p": lambda: prompt_cycle() + }[user_action]() + +get_user_action() #print(json.dumps(contents, indent="\t")) \ No newline at end of file diff --git a/main2.py b/main2.py deleted file mode 100644 index 60f948e..0000000 --- a/main2.py +++ /dev/null @@ -1,60 +0,0 @@ -import pymupdf -import os -import lib -import json - - -def retrieve_file_contents(path): - file_extraction_functions = { - "pdf": lambda path: lib.extract_pdf_content(path), - "jpg": lambda path: lib.extract_image_content(path), - "png": lambda path: lib.extract_image_content(path), - "txt": lambda path: lib.extract_pdf_content(path), - "mp3": lambda path: lib.extract_mp3_content(path), - } - - lib.read_files(path, files := []) - - contents = [] - - for file in files: - content = file_extraction_functions[file[0]](file[1]) - contents.append({ - "type": file[0], - "path": file[1], - "filename": file[2], - "content": content - }) - - return contents - - -def reload_files(): - path = input("Please provider path to folder: ") - pass - - -def add_files(): - path = input("Please provider path to folder: ") - pass - - -def prompt_cycle(): - while True: - prompt = input("Please enter prompt: ") - lib.prompt_embedding(prompt) - - -def get_user_action(): - user_action = input("Reload Files (r), Add Files (a), Prompt (p): ") - - { - "r": lambda: reload_files(), - "a": lambda: add_files(), - "p": lambda: prompt_cycle() - }[user_action]() - -# contents = retrieve_file_contents("C:\\SoftwareEng") -# print(json.dumps(contents, indent="\t")) - -get_user_action() \ No newline at end of file diff --git a/test_PDFparser.py b/test_PDFparser.py deleted file mode 100644 index fe11fe1..0000000 --- a/test_PDFparser.py +++ /dev/null @@ -1,13 +0,0 @@ -import pymupdf - -def extrahiere_text_als_string(pdf_datei): - doc = pymupdf.open(pdf_datei) - a = "" - for page in doc: - a += page.get_text() - return a - -pdf_datei = 'InformatikModulhandbuch.pdf' -a = extrahiere_text_als_string(pdf_datei) - -print(a) \ No newline at end of file diff --git a/test_get_files.py b/test_get_files.py deleted file mode 100644 index cde177c..0000000 --- a/test_get_files.py +++ /dev/null @@ -1,24 +0,0 @@ -import os - - -def read_files(path, output, filetypes=None): - print(os.path.join(path)) - if filetypes is None: - filetypes = ["pdf", "txt", "png", "jpg", "mp3"] - - for root, dirs, files in os.walk(path): - for file in files: - if (type := file.split(".")[-1]) in filetypes: - output.append((type, os.path.join(path, file))) - - for dir in dirs: - read_files(os.path.join(path, dir), output) - - break - - - -output = [] -read_files("C:\\SoftwareEng", output) - -print(output) \ No newline at end of file diff --git a/test_image.py b/test_image.py deleted file mode 100644 index a480c7b..0000000 --- a/test_image.py +++ /dev/null @@ -1,7 +0,0 @@ -import cv2 -import pytesseract - -def extract_image_content(path): - pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe' - img = cv2.imread(path) - return pytesseract.image_to_string(img) \ No newline at end of file diff --git a/test_mp3parser.py b/test_mp3parser.py deleted file mode 100644 index 2d7b02f..0000000 --- a/test_mp3parser.py +++ /dev/null @@ -1,12 +0,0 @@ -import whisper -import json - - -model = whisper.load_model('tiny') - - -result = model.transcribe(str("C:\\Users\\afist\\Downloads\\Aufzeichnung.mp3"), language='de', verbose=True) - -# Dump the results to a JSON file -with open('transcript.json', "w") as file: - json.dump(result['text'], file, indent=4) \ No newline at end of file