83 lines
3.3 KiB
Python
83 lines
3.3 KiB
Python
import os
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
|
|
import faiss
|
|
import numpy as np
|
|
import ollama
|
|
import gradio as gr
|
|
import os
|
|
import argparse
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
# === KONFIGURACJA ===
|
|
model_name = "gemma:2b" # Nazwa modelu Ollama
|
|
faiss_index_path = "faiss_index.idx" # Plik indeksu FAISS
|
|
kodeks_file = "./docs/kodekspracy.txt" # Plik z treścią kodeksu pracy
|
|
embedding_model = SentenceTransformer("all-MiniLM-L6-v2") # Model do embedowania tekstu
|
|
|
|
# === 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 = []
|
|
for article in articles:
|
|
if article.strip().startswith('Art.'):
|
|
documents.append(article.strip())
|
|
return 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
|
|
|
|
# === 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
|
|
|
|
# === 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ć."
|
|
|
|
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)}"
|
|
|
|
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.")
|
|
|
|
iface.launch(share=True) |