import os import time import pandas as pd import torch from datasets import Dataset, 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', 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') 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'], wrap=True), gr.Textbox(label='Время поиска, миллисекунды'), gr.Dataframe(label='Похожие связанные объекты', headers=['Название', 'Контекст', 'Сходство'], wrap=True), gr.Textbox(label='Время поиска, миллисекунды'), ], allow_flagging='never', examples=[ ['Функция', 'проектирование базы данных', 7, 'Должность', '', 7], ['Функция', 'написать руководство пользователя', 7, 'Должность', '', 7], ['Должность', 'программист', 12, 'Процесс', '', 7], ], title='Поиск по профстандартам', description='''Выберите тип объектов, который вы хотите найти, введите его название. Опционально укажите какие связанные объекты вы хотите найти.''', article=f'''

Поиск выполняется по реестру профессиональных стандартов минтруда.

В базе есть следующие данные:

Тип объектовКол-во
Предметные области{domains_ds.num_rows}
Процессы{generalized_functions_ds.num_rows}
Подпроцессы{particular_functions_ds.num_rows}
Функции{actions_ds.num_rows}
Должности{jobs_ds.num_rows}
Навыки{skills_ds.num_rows}
Знания{knowledges_ds.num_rows}

Для вычисления векторных представлений используется следующая модель:

Характеристика моделиЗначение
Модель{checkpoint}
Размер, Мб{round(get_model_size(model))}
Количество параметров, миллионы{round(model.num_parameters()/10**6)}
Размерность векторных представлений{get_embeddings('').shape[1]}
''', css='.w-full .col:nth-child(2) { flex-grow: 2 !important; }') ui.launch()