Create app.py

app.py

@@ -0,0 +1,154 @@
+import os
+import time
+
+import pandas as pd
+import torch
+
+from datasets import load_dataset
+from sentence_transformers import util
+from transformers import AutoTokenizer, AutoModel
+
+import gradio as gr
+
+device = torch.device('cpu')
+
+# Helpers
+def get_model_size(model):
+    param_size = 0
+    for param in model.parameters():
+        param_size += param.nelement() * param.element_size()
+    buffer_size = 0
+    for buffer in model.buffers():
+        buffer_size += buffer.nelement() * buffer.element_size()
+    return (param_size + buffer_size) / 1024**2
+
+# Load model
+checkpoint = 'sberbank-ai/sbert_large_mt_nlu_ru'
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModel.from_pretrained(checkpoint)
+model = model.to(device)
+
+def mean_pooling(token_embeddings, attention_mask):
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
+    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+    return sum_embeddings / sum_mask
+
+def get_embeddings(input):
+    encoded_input = tokenizer(input, padding=True, truncation=True, max_length=50, return_tensors='pt').to(device)
+    with torch.no_grad():
+        model_output = model(**encoded_input)
+    return mean_pooling(model_output[0], encoded_input['attention_mask']).cpu().numpy()
+
+# Load data
+ds_name = 'AresEkb/prof_standards_sbert_large_mt_nlu_ru'
+domains_ds = load_dataset(ds_name, 'domains').add_faiss_index(column='embeddings')
+generalized_functions_ds = load_dataset(ds_name, 'generalized_functions').add_faiss_index(column='embeddings')
+jobs_ds = load_dataset(ds_name, 'jobs').add_faiss_index(column='embeddings')
+particular_functions_ds = load_dataset(ds_name, 'particular_functions').add_faiss_index(column='embeddings')
+actions_ds = load_dataset(ds_name, 'actions').add_faiss_index(column='embeddings')
+skills_ds = load_dataset(ds_name, 'skills').add_faiss_index(column='embeddings')
+knowledges_ds = load_dataset(ds_name, 'knowledges').add_faiss_index(column='embeddings')
+
+indices = {'reg_number', 'generalized_function_id', 'particular_function_id'}
+
+entity_kinds = {
+    'Предметная область': domains_ds,
+    'Процесс': generalized_functions_ds,
+    'Подпроцесс': particular_functions_ds,
+    'Функция': actions_ds,
+    'Должность': jobs_ds,
+    'Навык': skills_ds,
+    'Знание': knowledges_ds,
+}
+
+# Main search logic
+def search(context_entity_kind, context_entity_name, context_entity_count,
+           target_entity_kind, target_entity_name, target_entity_count):
+    # Find similar context entities
+    start_time = time.perf_counter_ns()
+    context_ds = entity_kinds[context_entity_kind]
+    context_embedding = get_embeddings(context_entity_name)
+    scores, samples = context_ds.get_nearest_examples(
+        'embeddings', context_embedding, k=context_entity_count
+    )
+    cos_scores = util.cos_sim(context_embedding, samples['embeddings'])[0]
+    cos_scores = [round(x, 4) for x in cos_scores.tolist()]
+    results = pd.DataFrame({'name': samples['name'], 'score': cos_scores}).sort_values('score', ascending=False)
+    search_time = round((time.perf_counter_ns() - start_time) / 10**6)
+
+    # Get related entities
+    start_time = time.perf_counter_ns()
+    context_df = pd.DataFrame(samples).drop(columns=['embeddings']).rename(columns={'name': 'context_name'})
+    context_df['context_score'] = cos_scores
+    target_df = entity_kinds[target_entity_kind].to_pandas()
+    common_indices = list(indices.intersection(context_df.columns).intersection(target_df.columns))
+    target_ds = Dataset.from_pandas(context_df.merge(target_df, on=common_indices))
+    target_ds.add_faiss_index(column='embeddings')
+
+    # Find similar target entities
+    target_embedding = get_embeddings(target_entity_name)
+    scores, samples = target_ds.get_nearest_examples(
+        'embeddings', target_embedding, k=target_entity_count
+    )
+    cos_scores = util.cos_sim(target_embedding, samples['embeddings'])[0]
+    cos_scores = (cos_scores + torch.tensor(samples['context_score'])) / 2
+    cos_scores = [round(x, 4) for x in cos_scores.tolist()]
+    results2 = pd.DataFrame({'name': samples['name'], 'context_name': samples['context_name'], 'score': cos_scores}).sort_values('score', ascending=False)
+    search_time2 = round((time.perf_counter_ns() - start_time) / 10**6)
+
+    return [results.to_numpy(), search_time, results2.to_numpy(), search_time2]
+
+# User Interface
+ui = gr.Interface(
+    search,
+    [
+        gr.Radio(label='Тип объекта', choices=list(entity_kinds.keys()), value='Функция'),
+        gr.Textbox(label='Название объекта'),
+        gr.Slider(1, 20, 10, step=1, label='Кол-во объектов'),
+        gr.Radio(label='Тип связанного объекта', choices=list(entity_kinds.keys()), value='Должность'),
+        gr.Textbox(label='Название связанного объекта'),
+        gr.Slider(1, 20, 10, step=1, label='Кол-во связанных объектов'),
+    ],
+    [
+        gr.Dataframe(label='Похожие объекты', headers=['Название', 'Сходство'], datatype=['str', 'number']),
+        gr.Textbox(label='Время поиска, миллисекунды'),
+        gr.Dataframe(label='Похожие связанные объекты', headers=['Название', 'Контекст', 'Сходство']),
+        gr.Textbox(label='Время поиска, миллисекунды'),
+    ],
+    allow_flagging='never',
+    live=True,
+    examples=[
+        ['Функция', 'проектирование базы данных', 7, 'Должность', '', 7],
+        ['Функция', 'написать руководство пользователя', 7, 'Должность', '', 7],
+        ['Должность', 'программист', 12, 'Процесс', '', 7],
+    ],
+    title='Поиск по профстандартам',
+    description='''Выберите тип объектов, который вы хотите найти, введите его название.
+        Опционально укажите какие связанные объекты вы хотите найти.''',
+    article=f'''<p>Поиск выполняется по
+           <a href="https://profstandart.rosmintrud.ru/obshchiy-informatsionnyy-blok/natsionalnyy-reestr-professionalnykh-standartov/reestr-professionalnykh-standartov/">реестру</a>
+           профессиональных стандартов минтруда.</p>
+        <p>В базе есть следующие данные:</p>
+        <table>
+          <tr><th>Тип объектов</th><th>Кол-во</th></tr>
+          <tr><td>Предметные области</td><td>{domains_ds.num_rows}</td></tr>
+          <tr><td>Процессы</td><td>{generalized_functions_ds.num_rows}</td></tr>
+          <tr><td>Подпроцессы</td><td>{particular_functions_ds.num_rows}</td></tr>
+          <tr><td>Функции</td><td>{actions_ds.num_rows}</td></tr>
+          <tr><td>Должности</td><td>{jobs_ds.num_rows}</td></tr>
+          <tr><td>Навыки</td><td>{skills_ds.num_rows}</td></tr>
+          <tr><td>Знания</td><td>{knowledges_ds.num_rows}</td></tr>
+        </table>
+        <p>Для вычисления векторных представлений используется следующая модель:</p>
+        <table>
+          <tr><th>Характеристика модели</th><th>Значение</th></tr>
+          <tr><td>Модель</td><td><a href="https://huggingface.co/{checkpoint}">{checkpoint}</a></td></tr>
+          <tr><td>Размер, Мб</td><td>{round(get_model_size(model))}</td></tr>
+          <tr><td>Количество параметров, миллионы</td><td>{round(model.num_parameters()/10**6)}</td></tr>
+          <tr><td>Размерность векторных представлений</td><td>{get_embeddings('').shape[1]}</td></tr>
+        </table>
+    ''',
+    css='.w-full .col:nth-child(2) { flex-grow: 2 !important; }')
+
+ui.launch()