Update app.py

app.py

@@ -13,6 +13,7 @@ from docx import Document
 import asyncio
 import docx
 from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import SentenceTransformer
 import torch
 
 
@@ -32,6 +33,10 @@ api = HfApi()
 # Relativer Pfad zum Verzeichnis mit den Dokumenten
 DOCS_DIR = "kkg_dokumente"
 
+########################################################
+##########Ki Modell für Embeddings der Suchanfrage nutzen
+# Laden des Sentence-Transformer-Modells
+model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
 
 
 # Funktion zum Extrahieren von Text aus PDF
@@ -111,14 +116,16 @@ def download_link(doc_name):
     file_url = f"https://huggingface.co/spaces/alexkueck/kkg_suche/resolve/main/kkg_dokumente/{doc_name}?token={HF_READ}"
     return f'<a href="{file_url}" download="{doc_name}" target="_blank">{doc_name}</a>'
 
-# Laden des BERT-Modells und Tokenizers
-tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
-model = AutoModel.from_pretrained('bert-base-uncased')
+
 
 ###########################################################
 ############# KI um Suchanfrage zu Embedden ###############
 # Funktion zur Berechnung der Embeddings
 def get_embeddings(texts):
+    return model.encode(texts, convert_to_tensor=True)
+    
+# Funktion zur Berechnung der Embeddings
+def get_embeddingsback(texts):
     inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
@@ -136,8 +143,9 @@ def search_documents(query):
     # Embedding des Queries
     query_embedding = get_embeddings([query])
     
+    
     # Berechnung der Ähnlichkeit
-    similarities = cosine_similarity(query_embedding, text_embeddings).flatten()
+    similarities = cosine_similarity(query_embedding.cpu(), text_embeddings.cpu()).flatten()
 
     # Sortieren nach Relevanz
     related_docs_indices = similarities.argsort()[::-1]
@@ -150,33 +158,34 @@ def search_documents(query):
     cumulative_pages = [sum(num_pages_per_doc[:i+1]) for i in range(len(num_pages_per_doc))]
     
     for i in related_docs_indices:
-        doc_index = next(idx for idx, cumulative in enumerate(cumulative_pages) if i < cumulative)
-        page_index = i if doc_index == 0 else i - cumulative_pages[doc_index-1]
-        doc = documents[doc_index]
-        page = doc['pages'][page_index]
-        page_content = page['content']
-        header_content = page.get('header', '')
-            
-        # Überprüfen, ob der Suchtext in der Überschrift oder im Seiteninhalt enthalten ist
-        index_in_content = page_content.lower().find(query.lower())
-        index_in_header = header_content.lower().find(query.lower())
-            
-        if index_in_content != -1 or index_in_header != -1:
-            # Erstellen Sie einen Snippet für die Suchergebnisse
-            start = max(0, index_in_content - 400) if index_in_content != -1 else 0
-            end = min(len(page_content), index_in_content + 400) if index_in_content != -1 else len(page_content)
-            snippet = f"Aus <span class='doc-name'>{doc['file']}</span> (Seite <span class='page-number'>{page_index + 1}</span>):<br>"
-                
-            # Fügen Sie die Überschrift hinzu, falls vorhanden
-            if header_content:
-                snippet += f"<b>Überschrift:</b> {header_content}<br>"
+        if similarities[i] > 0:
+            doc_index = next(idx for idx, cumulative in enumerate(cumulative_pages) if i < cumulative)
+            page_index = i if doc_index == 0 else i - cumulative_pages[doc_index-1]
+            doc = documents[doc_index]
+            page = doc['pages'][page_index]
+            page_content = page['content']
+            header_content = page.get('header', '')
                 
-            snippet += f"...{page_content[start:end]}...<br><br>"
-            relevant_text += snippet
+            # Überprüfen, ob der Suchtext in der Überschrift oder im Seiteninhalt enthalten ist
+            index_in_content = page_content.lower().find(query.lower())
+            index_in_header = header_content.lower().find(query.lower())
                 
-            if doc['file'] not in relevant_docs:
-                relevant_docs[doc['file']] = []
-            relevant_docs[doc['file']].append(snippet)
+            if index_in_content != -1 or index_in_header != -1:
+                # Erstellen Sie einen Snippet für die Suchergebnisse
+                start = max(0, index_in_content - 400) if index_in_content != -1 else 0
+                end = min(len(page_content), index_in_content + 400) if index_in_content != -1 else len(page_content)
+                snippet = f"Aus <span class='doc-name'>{doc['file']}</span> (Seite <span class='page-number'>{page_index + 1}</span>):<br>"
+                    
+                # Fügen Sie die Überschrift hinzu, falls vorhanden
+                if header_content:
+                    snippet += f"<b>Überschrift:</b> {header_content}<br>"
+                    
+                snippet += f"...{page_content[start:end]}...<br><br>"
+                relevant_text += snippet
+                    
+                if doc['file'] not in relevant_docs:
+                    relevant_docs[doc['file']] = []
+                relevant_docs[doc['file']].append(snippet)
     
     results = list(relevant_docs.keys())
     return results, relevant_text