|
import os |
|
import streamlit as st |
|
import transformers |
|
import torch |
|
import tokenizers |
|
from typing import List, Dict |
|
|
|
st.subheader('Эта демонстрация позволяет поэксперементировать с моделями, которые оценивают, насколько предлагаемый ответ подходит к контексту диалога.') |
|
model_name = st.selectbox( |
|
'Выберите модель', |
|
('tinkoff-ai/response-quality-classifier-tiny', 'tinkoff-ai/response-quality-classifier-base', 'tinkoff-ai/response-quality-classifier-large') |
|
) |
|
auth_token = os.environ.get('TOKEN') or True |
|
|
|
@st.cache(hash_funcs={tokenizers.Tokenizer: lambda tokenizer: hash(tokenizer.to_str())}, allow_output_mutation=True) |
|
def load_model(model_name: str): |
|
with st.spinner('Loading models...'): |
|
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name, use_auth_token=auth_token) |
|
model = transformers.AutoModelForSequenceClassification.from_pretrained(model_name, use_auth_token=auth_token) |
|
if torch.cuda.is_available(): |
|
model = model.cuda() |
|
return tokenizer, model |
|
|
|
context_3 = 'привет' |
|
context_2 = 'привет!' |
|
context_1 = 'как дела?' |
|
|
|
st.markdown('👱🏻♀️ **Настя**: ' + context_3) |
|
st.markdown('🤖 **Диалоговый агент**: ' + context_2) |
|
st.markdown('👱🏻♀️ **Настя**: ' + context_1) |
|
response = st.text_input('🤖 Диалоговый агент:', 'норм') |
|
sample = { |
|
'context_3': context_3, |
|
'context_2': context_2, |
|
'context_1': context_1, |
|
'response': response |
|
} |
|
|
|
SEP_TOKEN = '[SEP]' |
|
CLS_TOKEN = '[CLS]' |
|
RESPONSE_TOKEN = '[RESPONSE_TOKEN]' |
|
MAX_SEQ_LENGTH = 128 |
|
sorted_dialog_columns = ['context_3', 'context_2', 'context_1', 'response'] |
|
|
|
|
|
def tokenize_dialog_data( |
|
tokenizer: transformers.PreTrainedTokenizer, |
|
sample: Dict, |
|
max_seq_length: int, |
|
sorted_dialog_columns: List, |
|
): |
|
""" |
|
Tokenize both contexts and response of dialog data separately |
|
""" |
|
len_message_history = len(sorted_dialog_columns) |
|
max_seq_length = min(max_seq_length, tokenizer.model_max_length) |
|
max_each_message_length = max_seq_length // len_message_history - 1 |
|
messages = [sample[k] for k in sorted_dialog_columns] |
|
result = {model_input_name: [] for model_input_name in tokenizer.model_input_names} |
|
messages = [str(message) if message is not None else '' for message in messages] |
|
tokens = tokenizer( |
|
messages, padding=False, max_length=max_each_message_length, truncation=True, add_special_tokens=False |
|
) |
|
for model_input_name in tokens.keys(): |
|
result[model_input_name].extend(tokens[model_input_name]) |
|
return result |
|
|
|
|
|
def merge_dialog_data( |
|
tokenizer: transformers.PreTrainedTokenizer, |
|
sample: Dict |
|
): |
|
cls_token = tokenizer(CLS_TOKEN, add_special_tokens=False) |
|
sep_token = tokenizer(SEP_TOKEN, add_special_tokens=False) |
|
response_token = tokenizer(RESPONSE_TOKEN, add_special_tokens=False) |
|
model_input_names = tokenizer.model_input_names |
|
result = {} |
|
for model_input_name in model_input_names: |
|
tokens = [] |
|
tokens.extend(cls_token[model_input_name]) |
|
for i, message in enumerate(sample[model_input_name]): |
|
tokens.extend(message) |
|
if i < len(sample[model_input_name]) - 2: |
|
tokens.extend(sep_token[model_input_name]) |
|
elif i == len(sample[model_input_name]) - 2: |
|
tokens.extend(response_token[model_input_name]) |
|
result[model_input_name] = torch.tensor([tokens]) |
|
if torch.cuda.is_available(): |
|
result[model_input_name] = result[model_input_name].cuda() |
|
return result |
|
|
|
|
|
@st.cache |
|
def inference(model_name: str, sample: dict): |
|
tokenizer, model = load_model(model_name) |
|
tokenized_dialog = tokenize_dialog_data(tokenizer, sample, MAX_SEQ_LENGTH, sorted_dialog_columns) |
|
tokens = merge_dialog_data(tokenizer, tokenized_dialog) |
|
with torch.inference_mode(): |
|
logits = model(**tokens).logits |
|
probas = torch.sigmoid(logits)[0].cpu().detach().numpy().tolist() |
|
return probas |
|
|
|
with st.spinner('Running inference...'): |
|
probas = inference(model_name, sample) |
|
st.metric( |
|
label='Вероятность того, что последний ответ диалогового агента релевантный', |
|
value=round(probas[0], 3) |
|
) |
|
st.metric( |
|
label='Вероятность того, что последний ответ диалогового агента вовлечённый', |
|
value=round(probas[1], 3) |
|
) |
|
|
