Update app.py

app.py

@@ -21,6 +21,8 @@ from utils import best_text_choice
 
 tokenizer = AutoTokenizer.from_pretrained("ai-forever/ru-en-RoSBERTa")
 model = AutoModel.from_pretrained("ai-forever/ru-en-RoSBERTa")
+# Load the DataFrame used in the best_text_choice function
+unique_sms_df = pd.read_parquet('unique_texts.parquet')
 
 MISTRAL_API_KEY = os.getenv('MISTRAL_API_KEY')
 token = os.getenv('GITHUB_TOKEN')
@@ -789,29 +791,60 @@ def generate_all_messages(desc, benefits, key_message, gender, generation, psych
     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()]
+    # Function to clean each message by removing metadata
+    def clean_message(msg):
+        # Remove metadata after '------'
+        if '------' in msg:
+            msg = msg.split('------')[0].strip()
+        return msg.strip()
     
-    # Ранжируем неперсонализированные сообщения
-    ranked_non_personalized_df = best_text_choice(non_personalized_list, unique_sms_df, tokenizer, model)
-    # Ранжируем персонализированные сообщения
-    ranked_personalized_df = best_text_choice(personalized_list, unique_sms_df, tokenizer, model)
+    # Split and clean non-personalized messages
+    non_personalized_list = [
+        clean_message(msg) 
+        for msg in non_personalized_messages.strip().split('\n\n') 
+        if msg.strip()
+    ]
+    
+    # Split and clean personalized messages
+    personalized_list = [
+        clean_message(msg) 
+        for msg in personalized_messages.strip().split('\n\n') 
+        if msg.strip()
+    ]
+    
+    # Rank non-personalized messages
+    ranked_non_personalized_df = best_text_choice(
+        non_personalized_list, unique_sms_df, tokenizer, model
+    )
     
-    # Предполагается, что функция best_text_choice возвращает DataFrame с колонкой 'text' или 'message'
-    # Извлекаем список сообщений из DataFrame
+    # Rank personalized messages
+    ranked_personalized_df = best_text_choice(
+        personalized_list, unique_sms_df, tokenizer, model
+    )
+    
+    # Extract messages from DataFrames
     ranked_non_personalized_messages_list = ranked_non_personalized_df['text'].tolist()
     ranked_personalized_messages_list = ranked_personalized_df['text'].tolist()
     
-    # Формируем строки для отображения
-    ranked_non_personalized_messages = '\n\n'.join(ranked_non_personalized_messages_list)
-    ranked_personalized_messages = '\n\n'.join(ranked_personalized_messages_list)
+    # Recompute lengths and reconstruct messages with metadata
+    def reconstruct_messages(messages_list):
+        display_list = []
+        for msg in messages_list:
+            msg_length = len(msg)
+            display_list.append(f"{msg}\n------\nКоличество знаков: {msg_length}")
+        return display_list
+    
+    ranked_non_personalized_display = reconstruct_messages(ranked_non_personalized_messages_list)
+    ranked_personalized_display = reconstruct_messages(ranked_personalized_messages_list)
+    
+    # Join the messages back into strings
+    ranked_non_personalized_messages = '\n\n'.join(ranked_non_personalized_display)
+    ranked_personalized_messages = '\n\n'.join(ranked_personalized_display)
     
     return ranked_non_personalized_messages, ranked_personalized_messages
 
+
 # ФУНКЦИИ ПРОВЕРОК (НАЧАЛО)
 
 # 1. Запрещенные слова