Update app.py

app.py

@@ -1,51 +1,49 @@
 import os
 import gradio as gr
-import spaces
 from datasets import load_dataset, concatenate_datasets
 from langchain.docstore.document import Document
 from langchain.vectorstores import FAISS
 from langchain_huggingface import HuggingFaceEmbeddings
 from tqdm import tqdm
-
-import torch
-print(torch.__version__)
-print(torch.cuda.is_available())
-
-# load env
 from dotenv import load_dotenv
+import pickle
 
-load_dotenv()
+# Импорты для перевода
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts.chat import ChatPromptTemplate
 
-key = os.getenv("HF_KEY")
+# Загружаем переменные окружения
+load_dotenv()
+hf_key = os.getenv("HF_KEY")
 
 # Путь для сохранения FAISS-индекса
 INDEX_PATH = "./faiss_index"
 
-# Инициализируем эмбеддинг-модель (используем модель из Hugging Face)
+# Инициализируем эмбеддинг-модель
 embeddings = HuggingFaceEmbeddings(model_name="fitlemon/bge-m3-uz-legal-matryoshka")
+translations = pickle.load(open("translations.pkl", "rb"))
 
 
 def update_faiss_index():
     """
-    Загружает датасет, преобразует данные в документы с метаданными,
+    Загружает датасеты, преобразует данные в документы с метаданными,
     создаёт FAISS-индекс и сохраняет его локально.
     """
-    # Загружаем датасет (например, сплит "train")
-    train_dataset = load_dataset("fitlemon/rag-labor-codex-dataset", token=key)["train"]
-    test_dataset = load_dataset("fitlemon/rag-labor-codex-dataset", token=key)["test"]
-    # combine train and test datasets
+    train_dataset = load_dataset("fitlemon/rag-labor-codex-dataset", token=hf_key)[
+        "train"
+    ]
+    test_dataset = load_dataset("fitlemon/rag-labor-codex-dataset", token=hf_key)[
+        "test"
+    ]
     dataset = concatenate_datasets([train_dataset, test_dataset])
-    # get rid off duplicate chunks
 
     docs = []
     unique_chunks = set()
     for row in tqdm(dataset, desc="Загрузка документов..."):
         chunk = row["chunk"]
-        # Если chunk уже добавлен, пропускаем его
         if chunk in unique_chunks:
             continue
         unique_chunks.add(chunk)
-
         doc = Document(
             page_content=chunk,
             metadata={
@@ -57,53 +55,112 @@ def update_faiss_index():
         docs.append(doc)
 
     print(f"Документы успешно загружены и преобразованы. Длина документов: {len(docs)}")
-    # Создаём FAISS-индекс на основе документов
     db = FAISS.from_documents(docs, embeddings)
-
-    # Сохраняем индекс в указанную директорию
     os.makedirs(INDEX_PATH, exist_ok=True)
     db.save_local(INDEX_PATH)
     print("FAISS индекс обновлён и сохранён в:", INDEX_PATH)
     return db
 
 
-# Если индекс ещё не создан, обновляем его, иначе загружаем существующий
 if not os.path.exists(INDEX_PATH):
     db = update_faiss_index()
 else:
     db = FAISS.load_local(INDEX_PATH, embeddings, allow_dangerous_deserialization=True)
     print("Загружен существующий FAISS индекс из:", INDEX_PATH)
 
-@spaces.GPU
-def retrieve_articles(query):
+
+def translate_ru_uz(message: str) -> str:
+    """
+    Переводит текст с русского на узбекский с использованием ChatOpenAI.
+    Пример: input: "отпуск" → output: "tatil".
     """
-    Принимает запрос пользователя, ищет в FAISS-индексе топ-3 наиболее релевантных документа
-    и возвращает отформатированный результат в Markdown.
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                "You are a helpful assistant that translates {input_language} to {output_language}. The subject of Text is Human Resources. Example input: отпуск. Output: tatil.",
+            ),
+            ("human", "{input}"),
+        ]
+    )
+    llm = ChatOpenAI(
+        api_key=os.environ["OPENAI_API_KEY"], model="gpt-4o-mini-2024-07-18"
+    )
+    chain = prompt | llm
+    response = chain.invoke(
+        {
+            "input_language": "Russian",
+            "output_language": "Uzbek",
+            "input": message,
+        }
+    )
+    return response.content
+
+
+def retrieve_articles(query, language):
     """
-    # Поиск по индексу: возвращает список из документов
-    results = db.similarity_search(query, k=3)
+    Если выбран язык "Russian", переводит запрос с русского на уз��екский.
+    Затем ищет в FAISS-индексе топ-3 наиболее релевантных документа и возвращает результат в Markdown.
+    """
+    if language == "Russian":
+        translated_query = translate_ru_uz(query)
+    else:
+        translated_query = query
+
+    results = db.similarity_search(translated_query, k=3)
 
-    # Форматируем результаты для вывода
     result_text = ""
     for doc in results:
         result_text += (
-            f"### Статья {doc.metadata['section']}: {doc.metadata['section_name']}\n"
+            f"### {doc.metadata['section']}: {doc.metadata['section_name']}\n"
         )
-        result_text += f"**Глава:** {doc.metadata['chapter_name']}\n\n"
-        result_text += f"**Текст статьи:**\n{doc.page_content}\n\n"
+        if language == "Russian":
+            result_text += f"**Текст статьи на русском:** {translations[doc.metadata['section']]}\n\n"
+        result_text += f"**Bo'lim:** {doc.metadata['chapter_name']}\n\n"
+        result_text += f"**Modda teksti:**\n{doc.page_content}\n\n"
         result_text += "---\n\n"
     return result_text
-    # return "### Hello World"
+    # return "Привет, мир!" if language == "Russian" else "Salom Dunyo!"
+
+
+def toggle_language(current_language: str) -> gr.update:
+    """
+    Переключает язык между "Russian" и "Uzbek".
+    """
+    new_language = "Uzbek" if current_language == "Russian" else "Russian"
+    return gr.update(value=new_language)
+
+
+# Создаём Gradio-интерфейс на основе Blocks
+with gr.Blocks() as demo:
+    gr.Markdown("# Поиск по Кодексу через Эмбеддинг Модель")
+    gr.Markdown(
+        "Введите ваш вопрос и выберите язык запроса. Если выбран русский, запрос будет переведен на узбекский перед поиском."
+    )
+
+    with gr.Row():
+        language_radio = gr.Radio(
+            choices=["Russian", "Uzbek"], label="Язык запроса", value="Russian"
+        )
+
+    with gr.Row():
+        query_input = gr.Textbox(
+            lines=3, placeholder="Введите ваш вопрос о кодексе...", label="Запрос"
+        )
+        search_button = gr.Button("Поиск")
 
+    output_markdown = gr.Markdown()
 
-# Создаём Gradio-интерфейс
-iface = gr.Interface(
-    fn=retrieve_articles,
-    inputs=gr.Textbox(lines=3, placeholder="Введите ваш вопрос о кодексе..."),
-    outputs=gr.Markdown(),
-    title="Поиск по Кодексу через FAISS",
-    description="Введите вопрос, и получите топ-3 наиболее релевантные статьи из кодекса.",
-)
+    search_button.click(
+        fn=retrieve_articles,
+        inputs=[query_input, language_radio],
+        outputs=output_markdown,
+    )
+    query_input.submit(
+        fn=retrieve_articles,
+        inputs=[query_input, language_radio],
+        outputs=output_markdown,
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()