Update func_ai.py

func_ai.py

@@ -1,73 +1,91 @@
+# func_ai.py
+
 import requests
 import torch
-# from googletrans import Translator
 from transformers import pipeline
 from deep_translator import GoogleTranslator
 import time
-import os 
+import os
+from datetime import datetime
+
 VECTOR_API_URL = os.getenv('API_URL')
 
-# translator = Translator()
+def log_message(message):
+    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+    print(f"[{timestamp}] {message}")
 
-sentiment_model = pipeline(
-    'sentiment-analysis',
-    model='cardiffnlp/twitter-xlm-roberta-base-sentiment',
-    tokenizer='cardiffnlp/twitter-xlm-roberta-base-sentiment',
-    device=0 if torch.cuda.is_available() else -1
-)
+# Инициализация моделей
+def init_models():
+    log_message("Инициализация моделей AI.")
+    sentiment_model = pipeline(
+        'sentiment-analysis',
+        model='cardiffnlp/twitter-xlm-roberta-base-sentiment',
+        tokenizer='cardiffnlp/twitter-xlm-roberta-base-sentiment',
+        device=0 if torch.cuda.is_available() else -1
+    )
+    
+    classifier = pipeline(
+        "zero-shot-classification",
+        model="valhalla/distilbart-mnli-12-6",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    return sentiment_model, classifier
 
-classifier = pipeline(
-    "zero-shot-classification",
-    model="valhalla/distilbart-mnli-12-6",
-    device=0 if torch.cuda.is_available() else -1
-)
+sentiment_model, classifier = init_models()
 
 def classify_comment(text):
     if not text:
-        print("Received empty text for classification.")
+        log_message("Получен пустой текст для классификации.")
         return "non-interrogative"
-    print(f"Classifying comment: {text}")
+    log_message(f"Классификация комментария: {text}")
     try:
         translated_text = GoogleTranslator(source='auto', target="en").translate(text)
-        print(f"Translated text: {translated_text}")
+        log_message(f"Переведенный текст: {translated_text}")
     except Exception as e:
-        print(f"Translation failed: {e}")
+        log_message(f"Ошибка при переводе: {e}")
         return "non-interrogative"
     if not translated_text:
-        print("Translation returned empty text.")
+        log_message("Перевод вернул пустой текст.")
         return "non-interrogative"
     
     try:
         result = classifier(translated_text, ["interrogative", "non-interrogative"], clean_up_tokenization_spaces=True)
-        print(f"Classification result: {result}")
+        log_message(f"Результат классификации: {result}")
     except Exception as e:
-        print(f"Classification failed: {e}")
+        log_message(f"Ошибка при классификации: {e}")
         return "non-interrogative"
     
     top_class = result['labels'][0]
-    print(f"Top class: {top_class}")
+    log_message(f"Верхний класс: {top_class}")
     return top_class
 
-def retrieve_from_vdb(query):
-    print(f"Отправка запроса к FastAPI сервису: {query}")
-    response = requests.post(f"{VECTOR_API_URL}/search/", json={"query": query})
-    if response.status_code == 200:
-        results = response.json().get("results", [])
-        print(f"Получено {len(results)} результатов: {results}")
-        return results
-    else:
-        print(f"Ошибка при поиске: {response.text}")
-        return []
-
 def analyze_sentiment(comments):
-    print("Начинаем анализ настроений.")
+    log_message("Начинаем анализ настроений.")
     results = []
     for i in range(0, len(comments), 50):
         batch = comments[i:i + 50]
-        print(f"Анализируем батч с {i} по {i + len(batch)} комментарий: {batch}")
-        batch_results = sentiment_model(batch)
-        print(f"Результаты батча: {batch_results}")
-        results.extend(batch_results)
+        log_message(f"Анализируем батч с {i} по {i + len(batch)} комментарий: {batch}")
+        try:
+            batch_results = sentiment_model(batch)
+            log_message(f"Результаты батча: {batch_results}")
+            results.extend(batch_results)
+        except Exception as e:
+            log_message(f"Ошибка при анализе настроений батча: {e}")
         time.sleep(1)  # Задержка для предотвращения перегрузки
-    print("Анализ настроений завершен. Общие результаты: {results}")
-    return results
+    log_message("Анализ настроений завершен.")
+    return results
+
+def retrieve_from_vdb(query):
+    log_message(f"Отправка запроса к FastAPI сервису: {query}")
+    try:
+        response = requests.post(f"{VECTOR_API_URL}/search/", json={"query": query})
+        if response.status_code == 200:
+            results = response.json().get("results", [])
+            log_message(f"Получено {len(results)} результатов: {results}")
+            return results
+        else:
+            log_message(f"Ошибка при поиске: {response.text}")
+            return []
+    except Exception as e:
+        log_message(f"Ошибка при запросе к векторной базе данных: {e}")
+        return []