From 51b505b6c1bec1eb1c55153702824210c531395e Mon Sep 17 00:00:00 2001
From: "l.gabrysiak" <l.gabrysiak@osadkowski.pl>
Date: Wed, 26 Feb 2025 15:56:21 +0100
Subject: [PATCH] mod

---
 gemma-faiss.py | 108 ++++++++++++++++++++++++++++---------------------
 1 file changed, 63 insertions(+), 45 deletions(-)

diff --git a/gemma-faiss.py b/gemma-faiss.py
index 6b61d53..8d7bba3 100644
--- a/gemma-faiss.py
+++ b/gemma-faiss.py
@@ -1,20 +1,23 @@
 import os
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-import numpy as np
 import faiss
-import gradio as gr
-import torch
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import numpy as np
 import ollama
+import gradio as gr
+import os
+import argparse
+from sentence_transformers import SentenceTransformer
 
-# 1️⃣ Inicjalizacja modelu do embeddingów
-embed_model = SentenceTransformer("all-MiniLM-L6-v2")
+# === KONFIGURACJA ===
+model_name = "gemma:2b"  # Nazwa modelu Ollama
+faiss_index_path = "faiss_index.idx"  # Plik indeksu FAISS
+kodeks_file = "kodeks_pracy.txt"  # Plik z treścią kodeksu pracy
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2")  # Model do embedowania tekstu
 
-# 2️⃣ Dodanie dokumentów i embeddingów
-def read_documents_from_file(file_path):
-    with open(file_path, 'r', encoding='utf-8') as file:
+# === KROK 1: WCZYTYWANIE KODEKSU PRACY ===
+def load_kodeks(filepath):
+    with open(filepath, 'r', encoding='utf-8') as file:
         content = file.read()
         articles = content.split('\n\n')
         documents = []
@@ -22,44 +25,59 @@ def read_documents_from_file(file_path):
             if article.strip().startswith('Art.'):
                 documents.append(article.strip())
     return documents
-#documents = [
-#    "Jak założyć firmę w Polsce?", 
-#    "Jak rozliczyć podatek VAT?", 
-#    "Procedura składania reklamacji w e-sklepie.",
-#    "Jakie dokumenty są potrzebne do rejestracji działalności?"
-#]
-file_path = './docs/kodekspracy.txt'  # Zmień na właściwą ścieżkę
-documents = read_documents_from_file(file_path)
 
-# 3️⃣ Wygenerowanie embeddingów
-embeddings = embed_model.encode(documents)
+# === KROK 2: TWORZENIE INDEKSU FAISS ===
+def create_faiss_index(sections):
+    embeddings = embedding_model.encode(sections, convert_to_numpy=True)  # Tworzenie wektorów
+    index = faiss.IndexFlatL2(embeddings.shape[1])  # Indeks FAISS
+    index.add(embeddings)  # Dodanie wektorów do FAISS
+    faiss.write_index(index, faiss_index_path)  # Zapis indeksu
+    return index, sections
 
-# 4️⃣ Inicjalizacja FAISS
-dim = embeddings.shape[1]  # Wymiar embeddingu
-index = faiss.IndexFlatL2(dim)
-index.add(np.array(embeddings, dtype=np.float32))
+# === KROK 3: WYSZUKIWANIE NAJBLIŻSZEGO FRAGMENTU ===
+def search_faiss(query, index, sections):
+    query_vector = embedding_model.encode([query], convert_to_numpy=True)  # Wektor zapytania
+    _, idx = index.search(query_vector, 1)  # Szukamy 1 najbliższego sąsiada
+    return sections[idx[0][0]]  # Zwracamy najbardziej pasujący fragment kodeksu
 
-# 5️⃣ Wczytanie modelu Ollama (Gemma 2)
-model_name = "hse.ably.do:latest"  # Ścieżka do modelu w systemie
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")
+# === KROK 4: GENEROWANIE ODPOWIEDZI Z OLLAMA ===
+def generate_response(user_query):
+    if not os.path.exists(faiss_index_path):
+        return "Błąd: Indeks FAISS nie istnieje. Uruchom aplikację z opcją --rebuild-index, aby go wygenerować."
 
-# 6️⃣ Funkcja wyszukiwania w FAISS
-def search(query, k=5):
-    query_embedding = embed_model.encode([query])  # Przekształć zapytanie w embedding
-    _, indices = index.search(np.array(query_embedding, dtype=np.float32), k)  # Szukaj w indeksie FAISS
-    return indices  # Zwróć indeksy najbardziej podobnych dokumentów
+    try:
+        index = faiss.read_index(faiss_index_path)  # Wczytanie indeksu FAISS
+    except Exception as e:
+        return f"Błąd: Nie można załadować indeksu FAISS. Spróbuj odbudować go za pomocą --rebuild-index. Szczegóły: {str(e)}"
 
-# 7️⃣ Funkcja generowania odpowiedzi z Ollama
-def generate_response(query):
-    indices = search(query)  # Znajdź najbardziej podobne dokumenty
-    relevant_documents = [documents[i] for i in indices[0]]  # Pobierz dokumenty na podstawie wyników wyszukiwania
+    sections = load_kodeks(kodeks_file)  # Wczytanie kodeksu
+    best_match = search_faiss(user_query, index, sections)  # Znalezienie najbliższego fragmentu
+
+    # Kontekst dla Ollama
+    prompt = f"Odpowiedz na pytanie na podstawie następującego tekstu:\n\n{best_match}\n\nPytanie: {user_query}"
+
+    response = ollama.chat(model=model_name, messages=[{"role": "user", "content": prompt}])
+    return response.get("message", "Błąd: Nie udało się wygenerować odpowiedzi.")
+
+# === KROK 5: URUCHOMIENIE INTERFEJSU WEBOWEGO ===
+iface = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(label="Zadaj pytanie o kodeks pracy"),
+    outputs=gr.Textbox(label="Odpowiedź"),
+    title="Asystent Kodeksu Pracy",
+    description="Wpisz pytanie, a system zwróci odpowiedni fragment kodeksu pracy."
+)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rebuild-index", action="store_true", help="Wymuś odbudowanie indeksu FAISS")
+    args = parser.parse_args()
+
+    if args.rebuild_index or not os.path.exists(faiss_index_path):
+        print("Tworzenie nowego indeksu FAISS...")
+        sections = load_kodeks(kodeks_file)
+        create_faiss_index(sections)
+    else:
+        print("Indeks FAISS już istnieje.")
     
-    prompt = " ".join(relevant_documents) + " " + query  # Przygotuj prompt
-    response = ollama.chat(model=model_name, messages=[{"role": "user", "content": prompt}])  # Generowanie odpowiedzi przez Ollama
-    
-    return response["text"]
-
-# 8️⃣ Interfejs Gradio (Open-WebUI)
-iface = gr.Interface(fn=generate_response, inputs="text", outputs="text")
-
-iface.launch(share=True)  # Uruchom interfejs
\ No newline at end of file
+    iface.launch(share=True)
\ No newline at end of file