Update app.py

app.py

@@ -5,7 +5,10 @@ from langchain_core.documents import Document  # Updated import
 from langchain_huggingface import HuggingFaceEmbeddings  # Updated import
 from langchain.evaluation import load_evaluator
 from langchain_community.vectorstores import Chroma
+from langchain_community.llms import HuggingFaceHub
+from langchain.prompts import ChatPromptTemplate
 from dotenv import load_dotenv
+import argparse
 import os
 import shutil  # Added import
 import numpy as np
@@ -17,6 +20,15 @@ load_dotenv()
 
 CHROMA_PATH = "chroma"
 DATA_PATH = ""  # Update this to your actual data path
+PROMPT_TEMPLATE = """
+Ответь на вопрос, используя только следующий контекст:
+
+{context}
+
+---
+
+Ответь на вопрос на основе приведенного контекста: {question}
+"""
 
 def main():
     # Создаем папки при необходимости
@@ -25,43 +37,48 @@ def main():
     
     generate_data_store()
 
-    # Инициализация модели
-    embedding_function = HuggingFaceEmbeddings(
+    help="Что означает Солнце на третьей ступени лестницы?"
+    process_query(help)
+
+def process_query(query_text: str):
+    # Инициализация эмбеддингов
+    embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-        cache_folder="model_cache",
-        model_kwargs={'device': 'cpu'},  # Форсируем использование CPU
-        encode_kwargs={'normalize_embeddings': True}
+        cache_folder="model_cache"
     )
-
-    # Тестирование эмбеддингов
-    test_embeddings(embedding_function)
-
-def test_embeddings(embeddings):
-    # Тест для "управитель"
-    vector = embeddings.embed_query("управитель")
-    print(f"Вектор для 'управитель' (первые 5 значений): {vector[:5]}")
-    print(f"Длина вектора: {len(vector)}")
-
-    # Инициализация эвалуатора
-    evaluator = load_evaluator(
-        "pairwise_embedding_distance",
-        embeddings=embeddings
+    
+    # Загрузка векторной БД
+    db = Chroma(
+        persist_directory=CHROMA_PATH, 
+        embedding_function=embeddings
     )
     
-    # Сравнение пар слов
-    pairs = [
-        ("управитель", "альмутен"),
-        ("Сатурн", "Венера")
-    ]
+    # Поиск по схожести
+    results = db.similarity_search_with_relevance_scores(query_text, k=3)
+    if not results or results[0][1] < 0.7:
+        print("Не найдено подходящих результатов.")
+        return
     
-    for word_a, word_b in pairs:
-        result = evaluator.evaluate_string_pairs(
-            prediction=word_a,
-            prediction_b=word_b
-        )
-        print(f"\nСравнение ({word_a}, {word_b}):")
-        print(f"Расстояние: {result['score']:.4f}")
-
+    # Формирование контекста
+    context_text = "\n\n---\n\n".join([doc.page_content for doc, _ in results])
+    
+    # Создание промпта
+    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
+    prompt = prompt_template.format(context=context_text, question=query_text)
+    
+    # Инициализация модели для генерации
+    model = HuggingFaceHub(
+        repo_id="google/flan-t5-small",
+        model_kwargs={"temperature": 0.5, "max_length": 512}
+    )
+    
+    # Генерация ответа
+    response_text = model.predict(prompt)
+    
+    # Форматирование вывода
+    sources = [doc.metadata.get("source", None) for doc, _ in results]
+    print(f"Ответ: {response_text}")
+    print(f"Источники: {sources}")
 
 def generate_data_store():
     documents = load_documents()