Create app.py

app.py

@@ -0,0 +1,146 @@
+import os
+import time
+
+import pandas as pd
+import torch
+
+from datasets import 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', split='train').add_faiss_index(column='embeddings')
+generalized_functions_ds = load_dataset(ds_name, 'generalized_functions', split='train').add_faiss_index(column='embeddings')
+jobs_ds = load_dataset(ds_name, 'jobs', split='train').add_faiss_index(column='embeddings')
+particular_functions_ds = load_dataset(ds_name, 'particular_functions', split='train').add_faiss_index(column='embeddings')
+actions_ds = load_dataset(ds_name, 'actions', split='train').add_faiss_index(column='embeddings')
+skills_ds = load_dataset(ds_name, 'skills', split='train').add_faiss_index(column='embeddings')
+knowledges_ds = load_dataset(ds_name, 'knowledges', split='train').add_faiss_index(column='embeddings')
+
+entity_kinds = {
+    'Предметная область': domains_ds,
+    'Процесс': generalized_functions_ds,
+    'Подпроцесс': particular_functions_ds,
+    'Функция': actions_ds,
+    'Должность': jobs_ds,
+    'Навык': skills_ds,
+    'Знание': knowledges_ds,
+}
+
+# Main logic
+def calc(entity_kind, entity1_name, entity2_name, entity3_name, entity_count):
+    start_time = time.perf_counter_ns()
+
+    if not entity1_name or not entity2_name and not entity3_name:
+        return [None, 0]
+
+    embedding1 = get_embeddings(entity1_name)
+    embedding2 = get_embeddings(entity2_name)
+    embedding3 = get_embeddings(entity3_name)
+
+    if entity2_name and entity3_name:
+        embedding = embedding1 - embedding2 + embedding3
+    elif entity2_name:
+        embedding = (embedding1 + embedding2) / 2
+    else:
+        embedding = (embedding1 + embedding3) / 2
+
+    context_ds = entity_kinds[entity_kind]
+    scores, samples = context_ds.get_nearest_examples(
+        'embeddings', embedding, k=entity_count
+    )
+    cos_scores = util.cos_sim(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)
+
+    return [results.to_numpy(), search_time]
+
+# User interface
+ui = gr.Interface(
+    calc,
+    [
+        gr.Radio(label='Тип искомого объекта', choices=list(entity_kinds.keys()), value='Функция'),
+        gr.Textbox(label='Исходный объект'),
+        gr.Textbox(label='Вычитаемый объект'),
+        gr.Textbox(label='Добавляемый объект'),
+        gr.Slider(1, 20, 10, step=1, label='Кол-во результатов'),
+    ],
+    [
+        gr.Dataframe(label='Результат вычисления', headers=['Название', 'Сходство'], datatype=['str', 'number'], wrap=True),
+        gr.Textbox(label='Время поиска, миллисекунды'),
+    ],
+    allow_flagging='never',
+    live=True,
+    examples=[
+        ['Функция', 'проектирование баз данных', 'база данных', 'интерфейс', 10],
+        ['Знание', 'управление проектами', 'проекты', 'производство', 10],
+        ['Функция', 'проектирование баз данных', 'проектирование баз данных', 'диагностика оборудования', 10],
+        ['Функция', 'проектирование автомобиля', 'управление автомобилем', 'управление данными', 10],
+        ['Функция', 'оценка данных', 'управление данными', 'управление персоналом', 10],
+        ['Функция', 'проектирование баз данных', 'диагностика оборудования', '', 10],
+    ],
+    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()