62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
import numpy as np
|
||
import faiss
|
||
import gradio as gr
|
||
import torch
|
||
from sentence_transformers import SentenceTransformer
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
import ollama
|
||
|
||
# 1️⃣ Inicjalizacja modelu do embeddingów
|
||
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
|
||
|
||
# 2️⃣ Dodanie dokumentów i embeddingów
|
||
def read_documents_from_file(file_path):
|
||
with open(file_path, '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
|
||
#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)
|
||
|
||
# 4️⃣ Inicjalizacja FAISS
|
||
dim = embeddings.shape[1] # Wymiar embeddingu
|
||
index = faiss.IndexFlatL2(dim)
|
||
index.add(np.array(embeddings, dtype=np.float32))
|
||
|
||
# 5️⃣ Wczytanie modelu Ollama (Gemma 2)
|
||
model_name = "./trained_model/gemma/Gemma-F16-LoRA.gguf" # Ścieżka do modelu w systemie
|
||
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b")
|
||
|
||
# 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
|
||
|
||
# 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
|
||
|
||
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() # Uruchom interfejs |