ably.do/gemma-faiss.py

93 lines
3.4 KiB
Python
Raw Normal View History

2025-02-26 09:35:02 -05:00
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
2025-02-26 09:30:23 -05:00
import faiss
2025-02-26 09:56:21 -05:00
import numpy as np
import ollama
2025-02-26 09:30:23 -05:00
import gradio as gr
2025-02-26 09:56:21 -05:00
import os
import argparse
2025-02-26 09:30:23 -05:00
from sentence_transformers import SentenceTransformer
2025-02-26 09:56:21 -05:00
# === KONFIGURACJA ===
2025-02-26 09:58:07 -05:00
model_name = "hse.ably.do:latest" # Nazwa modelu Ollama
2025-02-26 09:56:21 -05:00
faiss_index_path = "faiss_index.idx" # Plik indeksu FAISS
2025-02-26 10:04:44 -05:00
kodeks_file = "/home/ably.do/docs/kodekspracy.txt" # Plik z treścią kodeksu pracy
2025-02-26 09:56:21 -05:00
embedding_model = SentenceTransformer("all-MiniLM-L6-v2") # Model do embedowania tekstu
2025-02-26 09:30:23 -05:00
2025-02-26 09:56:21 -05:00
# === KROK 1: WCZYTYWANIE KODEKSU PRACY ===
def load_kodeks(filepath):
2025-02-26 10:04:44 -05:00
with open(filepath, "r", encoding="utf-8") as file:
2025-02-26 09:30:23 -05:00
content = file.read()
2025-02-26 10:04:44 -05:00
articles = content.split("\n\n") # Dzielimy na sekcje
return [article.strip() for article in articles if article.strip().startswith("Art.")]
2025-02-26 09:30:23 -05:00
2025-02-26 09:56:21 -05:00
# === 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
2025-02-26 09:30:23 -05:00
2025-02-26 09:56:21 -05:00
# === KROK 3: WYSZUKIWANIE NAJBLIŻSZEGO FRAGMENTU ===
2025-02-26 10:10:38 -05:00
def search_faiss(query, index, sections, top_k=3):
2025-02-26 10:04:44 -05:00
query_vector = embedding_model.encode([query], convert_to_numpy=True)
2025-02-26 10:10:38 -05:00
_, idx = index.search(query_vector, top_k) # Szukamy więcej wyników
results = [sections[i] for i in idx[0] if i < len(sections)]
return "\n\n".join(results) # Połącz kilka najlepszych fragmentów
2025-02-26 09:30:23 -05:00
2025-02-26 09:56:21 -05:00
# === KROK 4: GENEROWANIE ODPOWIEDZI Z OLLAMA ===
def generate_response(user_query):
if not os.path.exists(faiss_index_path):
2025-02-26 10:04:44 -05:00
return "Błąd: Indeks FAISS nie istnieje. Uruchom aplikację z opcją --rebuild-index."
2025-02-26 09:30:23 -05:00
2025-02-26 09:56:21 -05:00
try:
2025-02-26 10:04:44 -05:00
index = faiss.read_index(faiss_index_path)
2025-02-26 09:56:21 -05:00
except Exception as e:
2025-02-26 10:04:44 -05:00
return f"Błąd ładowania FAISS: {str(e)}"
2025-02-26 09:30:23 -05:00
2025-02-26 10:04:44 -05:00
sections = load_kodeks(kodeks_file)
best_match = search_faiss(user_query, index, sections)
2025-02-26 09:56:21 -05:00
2025-02-26 10:08:36 -05:00
# 👀 DEBUG: Sprawdź, co zwraca FAISS
print(f"🔍 Najlepsze dopasowanie FAISS dla '{user_query}':\n{best_match}")
prompt = f"""
Odpowiedz na pytanie na podstawie następującego tekstu:
{best_match}
Pytanie: {user_query}
Podaj dokładny tekst artykułu, jeśli go znajdziesz w treści powyżej.
"""
2025-02-26 09:56:21 -05:00
response = ollama.chat(model=model_name, messages=[{"role": "user", "content": prompt}])
2025-02-26 09:30:23 -05:00
2025-02-26 10:08:36 -05:00
print(f"📝 Odpowiedź modelu:\n{response}") # 👀 DEBUG: Sprawdź odpowiedź Ollama
2025-02-26 10:04:44 -05:00
return response.get("message", response.get("content", "Błąd: Nie udało się wygenerować odpowiedzi."))
# === KROK 5: INTERFEJS WEBOWY ===
2025-02-26 09:56:21 -05:00
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."
)
2025-02-26 09:30:23 -05:00
2025-02-26 09:56:21 -05:00
if __name__ == "__main__":
parser = argparse.ArgumentParser()
2025-02-26 10:04:44 -05:00
parser.add_argument("--rebuild-index", action="store_true", help="Odbudowanie indeksu FAISS")
2025-02-26 09:56:21 -05:00
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)