Update app.py

app.py

@@ -16,6 +16,11 @@ from collections import defaultdict
 import requests
 import base64
 import io
+from transformers import AutoTokenizer, AutoModel
+from utils import best_text_choice
+
+tokenizer = AutoTokenizer.from_pretrained("ai-forever/ru-en-RoSBERTa")
+model = AutoModel.from_pretrained("ai-forever/ru-en-RoSBERTa")
 
 MISTRAL_API_KEY = os.getenv('MISTRAL_API_KEY')
 token = os.getenv('GITHUB_TOKEN')
@@ -783,6 +788,24 @@ def generate_all_messages(desc, benefits, key_message, gender, generation, psych
         time.sleep(1)
     save_statistics_to_github(approach_stats)
 
+def rank_messages(non_personalized_messages, personalized_messages):
+    # Предполагается, что у вас есть DataFrame unique_sms_df, используемый в функции best_text_choice
+    unique_sms_df = pd.read_parquet('unique_texts.parquet')
+    
+    # Разделяем сообщения на отдельные строки
+    non_personalized_list = [msg.strip() for msg in non_personalized_messages.strip().split('\n\n') if msg.strip()]
+    personalized_list = [msg.strip() for msg in personalized_messages.strip().split('\n\n') if msg.strip()]
+    
+    # Ранжируем неперсонализированные сообщения
+    ranked_non_personalized = best_text_choice(non_personalized_list, unique_sms_df, tokenizer, model)
+    # Ранжируем персонализированные сообщения
+    ranked_personalized = best_text_choice(personalized_list, unique_sms_df, tokenizer, model)
+    
+    # Формируем строки для отображения
+    ranked_non_personalized_messages = '\n\n'.join(ranked_non_personalized)
+    ranked_personalized_messages = '\n\n'.join(ranked_personalized)
+    
+    return ranked_non_personalized_messages, ranked_personalized_messages
 
 # ФУНКЦИИ ПРОВЕРОК (НАЧАЛО)
 
@@ -1763,6 +1786,7 @@ with gr.Blocks() as demo:
             with gr.Row():
                 non_personalized_messages = gr.Textbox(label="Стандартные сообщения", lines=12, interactive=False)
                 personalized_messages = gr.Textbox(label="Персонализированные сообщения", lines=12, interactive=False)
+                rank_button = gr.Button("Ранжировать")
 
             # Сначала переключаем вкладку, потом запускаем генерацию сообщений
             btn_to_prompts.click(
@@ -1779,4 +1803,10 @@ with gr.Blocks() as demo:
                 ]
             )
 
+            rank_button.click(
+                fn=rank_messages,
+                inputs=[non_personalized_messages, personalized_messages],
+                outputs=[non_personalized_messages, personalized_messages]
+            )
+
 demo.launch()