fix out

app.py

@@ -15,7 +15,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def main():
-    st.subheader("CAIF")
+    st.header("CAIF")
     cls_model_name = st.selectbox(
         'Выберите модель классификации',
         ('tinkoff-ai/response-quality-classifier-tiny', 'tinkoff-ai/response-quality-classifier-base',
@@ -26,6 +26,7 @@ def main():
         ('sberbank-ai/rugpt3small_based_on_gpt2',)
     )
     prompt = st.text_input("Начало текста:", "Привет")
+    alpha = st.slider("Alpha:", min_value=-10, max_value=10, step=1)
     auth_token = os.environ.get('TOKEN') or True
     with st.spinner('Running inference...'):
         text = inference(lm_model_name=lm_model_name, cls_model_name=cls_model_name, prompt=prompt)
@@ -49,7 +50,7 @@ def load_sampler(cls_model_name, lm_tokenizer):
 
 
 @st.cache
-def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool = True) -> str:
+def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool = True, alpha: float = 5) -> str:
     generator = load_generator(lm_model_name=lm_model_name)
     lm_tokenizer = transformers.AutoTokenizer.from_pretrained(lm_model_name)
     caif_sampler = load_sampler(cls_model_name=cls_model_name, lm_tokenizer=lm_tokenizer)
@@ -59,7 +60,7 @@ def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool =
         "top_k": 20,
         "temperature": 1.0,
         "top_k_classifier": 100,
-        "classifier_weight": 5,
+        "classifier_weight": alpha,
     }
     generator.set_ordinary_sampler(ordinary_sampler)
     if device == "cpu":