Update app.py

app.py

@@ -27,7 +27,7 @@ def load_custom_model():
 # Функция для генерации текста
 def generate_text(model, tokenizer, prompt, max_len, temperature):
     input_ids = tokenizer.encode(prompt, return_tensors='pt')
-
+    attention_mask = (input_ids != tokenizer.pad_token_id).long()
     # Генерация текста
     output = model.generate(
         input_ids,
@@ -38,6 +38,7 @@ def generate_text(model, tokenizer, prompt, max_len, temperature):
         repetition_penalty=1.2,       # Штраф за повторение слов или фраз
         no_repeat_ngram_size=4,       # Запрет на повторение n-грамм (например, биграмм)
         do_sample=True,                # Включение сэмплинга для большей разнообразности
+        attention_mask=attention_mask
         pad_token_id=tokenizer.eos_token_id
     )
 
@@ -64,12 +65,12 @@ def main():
 
     # Блок для генерации текста с кастомной моделью
     st.subheader("Custom Model Text Generation")
-    prompt_custom = st.text_area("Введите фразу для генерации с кастомной моделью:", value="Когда-то давно")
-    generate_button_custom = st.button("Сгенерировать текст с кастомной моделью")
+    prompt_custom = st.text_area("Enter a phrase to generate with the updated model:", value="Когда-то давно")
+    generate_button_custom = st.button("Generate!")
 
     if generate_button_custom:
         generated_text_custom = generate_text(model_custom, tokenizer_custom, prompt_custom, max_len, temperature)
-        st.subheader("Результат генерации с кастомной моделью:")
+        st.subheader("Result:")
         st.write(generated_text_custom)
 
 if __name__ == "__main__":