Update README.md

README.md

@@ -32,81 +32,15 @@ The preferable usage:
 
 ```python
 # pip install transformers
-import transformers
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-from typing import List, Dict
-tokenizer = AutoTokenizer.from_pretrained("tinkoff-ai/response-quality-classifier-tiny")
-model = AutoModel.from_pretrained("tinkoff-ai/response-quality-classifier-tiny")
+tokenizer = AutoTokenizer.from_pretrained("/mnt/chatbot_models2/chit-chat/experiments/crossencoder_hf/rubert-base-sentence/dialogs_whole")
+model = AutoModelForSequenceClassification.from_pretrained("/mnt/chatbot_models2/chit-chat/experiments/crossencoder_hf/rubert-base-sentence/dialogs_whole")
 # model.cuda()
-context_3 = 'привет'
-context_2 = 'привет!'
-context_1 = 'как дела?'
-response = 'у меня все хорошо, а у тебя как?'
-
-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
-
-tokenized_dialog = tokenize_dialog_data(tokenizer, sample, MAX_SEQ_LENGTH, sorted_dialog_columns)
-tokens = merge_dialog_data(tokenizer, tokenized_dialog)
+inputs = tokenizer('привет[SEP]привет![SEP]как дела?[RESPONSE_TOKEN]норм',
+                   padding=True, max_length=128, truncation=True, add_special_tokens=False, return_tensors='pt')
 with torch.inference_mode():
-    logits = model(**tokens).logits
+    logits = model(**inputs).logits
     probas = torch.sigmoid(logits)[0].cpu().detach().numpy()
-
 print(probas)
 ```