83 lines
2.0 KiB
Python
83 lines
2.0 KiB
Python
import ollama
|
|
import pymupdf
|
|
import os
|
|
import lib
|
|
import json
|
|
import psycopg2
|
|
|
|
|
|
|
|
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
|
|
|
|
contents = retrieve_file_contents("C:\\Users\\afist\\Downloads\\SoftwareEng")
|
|
|
|
|
|
embedding_list = []
|
|
|
|
def create_embeddings(pContent):
|
|
conn = psycopg2.connect(
|
|
dbname="embeddings",
|
|
user="python",
|
|
password="PasswordPassword123",
|
|
host="localhost",
|
|
port="5555"
|
|
)
|
|
cur = conn.cursor()
|
|
|
|
create_table_query = '''
|
|
create table if not exists dbtable (
|
|
id SERIAL PRIMARY KEY,
|
|
filename TEXT NOT NULL,
|
|
embedding VECTOR NOT NULL
|
|
)
|
|
'''
|
|
|
|
|
|
cur.execute(create_table_query)
|
|
|
|
cur.execute('CREATE EXTENSION IF NOT EXISTS vector;')
|
|
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)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
create_embeddings(contents)
|
|
|
|
|
|
#print(json.dumps(contents, indent="\t")) |