good

app.py

@@ -10,10 +10,15 @@ 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
-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()
+
+@st.cache
+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 = 'привет!'
@@ -83,12 +88,20 @@ def merge_dialog_data(
            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)
-with torch.inference_mode():
-    logits = model(**tokens).logits
-    probas = torch.sigmoid(logits)[0].cpu().detach().numpy()
 
+@st.cache
+def inference(model_name: str, sample: dict):
+    tokenizer, model = load_model(model_name)
+    with st.spinner('Running inference...'):
+        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()
+        return probas
+
+
+probas = inference(model_name, sample)
 st.metric(
     label='Вероятность того, что последний ответ диалогового агента релевантный',
     value=probas[0]