112 lines
3.2 KiB
Python
112 lines
3.2 KiB
Python
import os, pymupdf, whisper, json
|
|
|
|
import cv2
|
|
import psycopg2
|
|
import pymongo
|
|
import pytesseract
|
|
import ollama
|
|
import io
|
|
from PIL import Image
|
|
from aiohttp.web_response import json_response
|
|
|
|
import ollama_conn
|
|
import redis_conn
|
|
|
|
|
|
def read_files(path, output, filetypes=None):
|
|
if filetypes is None:
|
|
filetypes = ["pdf", "txt", "png", "jpg", "mp3"]
|
|
|
|
for root, dirs, files in os.walk(path):
|
|
for file in files:
|
|
if (file_type := file.split(".")[-1]) in filetypes:
|
|
output.append((file_type, os.path.join(path, file), file))
|
|
|
|
for folder in dirs:
|
|
read_files(os.path.join(path, folder), output)
|
|
|
|
break
|
|
|
|
def extract_pdf_content(pdf_datei):
|
|
doc = pymupdf.open(pdf_datei)
|
|
a = ""
|
|
for page in doc:
|
|
a += page.get_text()
|
|
return a
|
|
|
|
|
|
def extract_mp3_content(mp3_datei):
|
|
model = whisper.load_model('tiny')
|
|
|
|
result = model.transcribe(str(mp3_datei), language='de', verbose=True)
|
|
|
|
# with open('transcript.json', "w") as f:
|
|
# json.dump(result['text'], f, indent=4)
|
|
|
|
return result["text"]
|
|
|
|
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)
|
|
|
|
def extract_text_and_pictures(pdf_datei):
|
|
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
|
|
doc = pymupdf.open(pdf_datei)
|
|
text = ""
|
|
for seite in doc:
|
|
text += seite.get_text()
|
|
for bild in seite.get_images(full=True):
|
|
xref = bild[0]
|
|
basis_bild = doc.extract_image(xref)
|
|
bild_daten = basis_bild["image"]
|
|
bild_objekt = Image.open(io.BytesIO(bild_daten))
|
|
bild_text = pytesseract.image_to_string(bild_objekt, lang="deu")
|
|
text += " " + bild_text
|
|
return text
|
|
|
|
def prompt_embedding(prompt):
|
|
embedding = ollama.embeddings(model="mxbai-embed-large", prompt=prompt)
|
|
|
|
conn = psycopg2.connect(
|
|
dbname="embeddings",
|
|
user="python",
|
|
password="PasswordPassword123",
|
|
host="localhost",
|
|
port="5555"
|
|
)
|
|
cur = conn.cursor()
|
|
|
|
cur.execute(f"select id, filepath, embedding <-> %s::vector as distance from dbtable order by distance limit 20;", (embedding["embedding"],))
|
|
|
|
result = cur.fetchall()
|
|
|
|
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)
|
|
|
|
ollama_response = ollama_conn.ask_ollama(prompt, res)
|
|
print(ollama_response)
|
|
redis_conn.cache_prompt_to_redis(prompt.lower().strip(), ollama_response)
|
|
|
|
conn.close()
|
|
|
|
def get_data_from_mongo(dataset):
|
|
result_string = ""
|
|
|
|
client = pymongo.MongoClient('mongodb://python:PasswordPassword123@localhost:27017/')
|
|
mongodb = client['document_table']
|
|
collection = mongodb['documents']
|
|
|
|
for set in dataset:
|
|
json_data = collection.find_one({"doc_id": set[0]})
|
|
result_string += "Dateipfad: " + json_data["filepath"] + "; Inhalt: " + json_data["chunk_content"] + "; "
|
|
|
|
return result_string |