First commit

app.py

@@ -0,0 +1,43 @@
+
+import gradio as gr
+from dotenv import load_dotenv
+from search import search_bm25, search_exact, prepare_data, merge_results
+import os
+
+load_dotenv()
+
+data = prepare_data()
+
+HF_TOKEN = os.getenv('HF_TOKEN')
+hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "budu_search_data_new")
+
+def search_handler(query: str):
+    results, exact_results = (
+        search_bm25(query, data),
+        search_exact(query, data)
+    )
+    
+    json_results = merge_results(exact_results, results)
+    return {'results': json_results}
+
+def create_ui(query):
+    recommendations = []
+    results = search_handler(query)
+    for result in results['results'][:5]:
+        recommendations.append(f"<div style='padding: 10px; border-bottom: 1px solid #ddd;'>{result['name']}</div>")
+    
+    return gr.HTML(f"<div style='max-height: 400px; overflow-y: auto;'>{''.join(recommendations)}</div>")
+
+iface = gr.Interface(
+    fn=create_ui,
+    inputs=gr.Textbox(label="Введите запрос"),
+    outputs=gr.HTML(),  # Use HTML to render custom styled output,
+    allow_flagging='manual',
+    flagging_callback = gr.CSVLogger(),
+    flagging_options = ['Хорошая рекомендаация',
+                        'Плохая рекомендаация'],
+    title="Поисковая система BUDU",
+)
+
+iface.launch()
+

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+python-dotenv
+rank_bm25
+transformers
+pandas 
+psycopg2-binary

search.py

@@ -0,0 +1,133 @@
+from transformers import AutoTokenizer
+from rank_bm25 import BM25Okapi
+import json
+import re
+
+tokenizer = AutoTokenizer.from_pretrained("deepvk/USER-bge-m3")
+
+
+with open('search.json', 'r', encoding='utf-8') as file:
+    data = json.load(file)
+
+
+def prepare_data(data=data):
+    for item in data:
+        
+        corpus = [item['name'].lower()] + [x.lower() for x in item['synonyms']]
+        corpus = " ".join(corpus)
+        
+        item['corpus'] = corpus 
+        item['tokenized_corpus'] = tokenizer.tokenize(corpus)
+    
+    return data
+
+
+def clean_text(text):
+    """
+    Очищает текст от специальных символов, чтобы можно было искать по содержимому.
+    """
+    return re.sub(r'[^\w\s]', '', text).lower()
+
+def search_bm25(query, data):
+    """
+    Выполняет поиск по запросу с использованием BM25 + токенизатора.
+
+    query: строка, поисковый запрос.
+    data: список словарей, содержащих информацию о каждом элементе (name, id, synonyms, tokenized_corpus).
+    
+    return: список словарей с уникальными name, id и score.
+    """
+
+    tokenized_corpus = [item['tokenized_corpus'] for item in data]
+    bm25 = BM25Okapi(tokenized_corpus)
+
+    tokenized_query = tokenizer.tokenize(query)
+    doc_scores = bm25.get_scores(tokenized_query)
+    
+
+    results = []
+    for idx in reversed(doc_scores.argsort()):
+        item = data[idx]
+        if item['name'] not in [result['name'] for result in results] and doc_scores[idx] > 0:
+            results.append({
+                'name': item['name'],
+                'id': item['id'],
+                'score': float(doc_scores[idx])
+            })
+    
+    return results
+
+def search_exact(query, data):
+    """
+    Выполняет точный поиск подстроки в 'name' и 'synonyms'. 
+    Слово либо является началом слова, либо целым словом.
+    
+    """
+    results = []
+    cleaned_query = clean_text(query)
+
+
+    query_regex = re.compile(r'\b' + re.escape(cleaned_query) + r'\b|\b' + re.escape(cleaned_query))  # Выбираем только целые слова или начала слов (чтобы не искать подстроки)
+    # Сначала ищем подстроку в 'name'
+    for item in data:
+        cleaned_name = clean_text(item['name'])
+        if query_regex.search(cleaned_name):
+            results.append({
+                'name': item['name'],
+                'id': item['id'],
+                'score': 1  # В данном случае score не важен, но можно оставить
+            })
+
+    # Затем ищем подстроку в 'synonyms'
+    for item in data:
+        for synonym in item['synonyms']:
+            cleaned_synonym = clean_text(synonym)
+            if query_regex.search(cleaned_synonym):
+                if not any(res['id'] == item['id'] for res in results):  # Избегаем дублирования
+                    results.append({
+                        'name': item['name'],
+                        'id': item['id'],
+                        'score': 1
+                    })
+                break  # Достаточно найти одно совпадение среди синонимов
+
+    return results
+
+#TODO: Добавить поиск подкатегорий
+
+
+def merge_results(results1, results2):
+    """
+    Объединяет два списка результатов с чередованием элементов, избегая дубликатов.
+    
+    results1: первый список результатов (список словарей с ключами 'name', 'id', 'score')
+    results2: второй список результатов (список словарей с ключами 'name', 'id', 'score')
+    
+    return: объединённый список результатов
+    """
+    merged_results = []
+    seen_ids = set()  # Набор для отслеживания уникальных id
+
+    # Чередуем элементы из двух списков
+    for res1, res2 in zip(results1, results2):
+        # Добавляем элемент из первого списка, если его id ещё не было
+        if res1['id'] not in seen_ids:
+            merged_results.append(res1)
+            seen_ids.add(res1['id'])
+
+        # Добавляем элемент из второго списка, если его id ещё не было
+        if res2['id'] not in seen_ids:
+            merged_results.append(res2)
+            seen_ids.add(res2['id'])
+
+    for res in results1[len(results2):]:
+        if res['id'] not in seen_ids:
+            merged_results.append(res)
+            seen_ids.add(res['id'])
+
+    for res in results2[len(results1):]:
+        if res['id'] not in seen_ids:
+            merged_results.append(res)
+            seen_ids.add(res['id'])
+
+    return merged_results