add languages

app.py

@@ -15,35 +15,76 @@ from generator import Generator
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+ATTRIBUTE_MODELS = {
+    "Russian": (
+        "cointegrated/rubert-tiny-toxicity",
+        'tinkoff-ai/response-quality-classifier-tiny',
+        'tinkoff-ai/response-quality-classifier-base',
+        'tinkoff-ai/response-quality-classifier-large',
+        "SkolkovoInstitute/roberta_toxicity_classifier",
+        "SkolkovoInstitute/russian_toxicity_classifier"
+    ),
+    "English": (
+        "unitary/toxic-bert",
+    )
+}
+
+LANGUAGE_MODELS = {
+    "Russian": ('sberbank-ai/rugpt3small_based_on_gpt2',),
+    "Eanglish": ("distilgpt2")
+}
+
+ATTRIBUTE_MODEL_LABEL = {
+    "Russian": 'Выберите модель классификации',
+    "English": "Choose attribute model"
+}
+
+LM_LABEL = {
+    "English": "Choose language model",
+    "Russian": "Выберите языковую модель"
+}
+
+ATTRIBUTE_LABEL = {
+    "Russian": "Веберите нужный атрибут текста",
+    "English": "Choose desired attribute",
+}
+
+TEXT_PROMPT_LABEL = {
+    "English": "Text prompt",
+    "Russian": "Начало текста"
+}
+
+PROMPT_EXAMPLE = {
+    "English": "Hello, today I",
+    "Russian": "Привет, сегодня я"
+}
+
 
 def main():
     st.header("CAIF")
+    language = st.selectbox("Language", ("English", "Russian"))
     cls_model_name = st.selectbox(
-        'Выберите модель классификации',
-        (
-            'tinkoff-ai/response-quality-classifier-tiny',
-            'tinkoff-ai/response-quality-classifier-base',
-            'tinkoff-ai/response-quality-classifier-large',
-            "SkolkovoInstitute/roberta_toxicity_classifier"
-        )
+        ATTRIBUTE_MODEL_LABEL[language],
+        ATTRIBUTE_MODELS[language]
+
     )
     lm_model_name = st.selectbox(
-        'Выберите языковую модель',
-        ('sberbank-ai/rugpt3small_based_on_gpt2',)
+        LM_LABEL[language],
+        LANGUAGE_MODELS[language]
     )
     cls_model_config = AutoConfig.from_pretrained(cls_model_name)
     if cls_model_config.problem_type == "multi_label_classification":
         label2id = cls_model_config.label2id
-        label_key = st.selectbox("Веберите нужный атрибут текста", label2id.keys())
+        label_key = st.selectbox(ATTRIBUTE_LABEL[language], label2id.keys())
         target_label_id = label2id[label_key]
     else:
         label2id = cls_model_config.label2id
         print(list(label2id.keys()))
-        label_key = st.selectbox("Веберите нужный атрибут текста", [list(label2id.keys())[-1]])
-        target_label_id = 1
-    prompt = st.text_input("Начало текста:", "Привет")
-    alpha = st.slider("Alpha:", min_value=-10, max_value=10, step=1, value=0)
-    entropy_threshold = st.slider("Entropy Threshold:", min_value=0., max_value=5., step=.1, value=0.)
+        label_key = st.selectbox(ATTRIBUTE_LABEL[language], [list(label2id.keys())[-1]])
+        target_label_id = 0
+    prompt = st.text_input(TEXT_PROMPT_LABEL[language], PROMPT_EXAMPLE[language])
+    alpha = st.slider("Alpha", min_value=-10, max_value=10, step=1, value=0)
+    entropy_threshold = st.slider("Entropy threshold", min_value=0., max_value=5., step=.1, value=0.)
     auth_token = os.environ.get('TOKEN') or True
     with st.spinner('Running inference...'):
         text = inference(