Update func_ai.py

func_ai.py

@@ -1,18 +1,17 @@
 import requests
 import torch
-# from googletrans import Translator
 from transformers import pipeline
 from deep_translator import GoogleTranslator
 import time
 import os 
-VECTOR_API_URL = os.getenv('API_URL')
 
-# translator = Translator()
+VECTOR_API_URL = os.getenv('API_URL')
 
+# Initialize the sentiment analysis model with distilbert
 sentiment_model = pipeline(
     'sentiment-analysis',
-    model='cardiffnlp/twitter-xlm-roberta-base-sentiment',
-    tokenizer='cardiffnlp/twitter-xlm-roberta-base-sentiment',
+    model='distilbert/distilbert-base-uncased-finetuned-sst-2-english',
+    tokenizer='distilbert/distilbert-base-uncased-finetuned-sst-2-english',
     device=0 if torch.cuda.is_available() else -1
 )
 
@@ -65,9 +64,23 @@ def analyze_sentiment(comments):
     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)
+
+        # Translate each comment in the batch
+        try:
+            translated_batch = [GoogleTranslator(source='auto', target='en').translate(comment) for comment in batch]
+            print(f"Translated batch: {translated_batch}")
+        except Exception as e:
+            print(f"Translation failed: {e}")
+            continue
+
+        # Analyze sentiment on the translated batch
+        try:
+            batch_results = sentiment_model(translated_batch)
+            print(f"Результаты батча: {batch_results}")
+            results.extend(batch_results)
+        except Exception as e:
+            print(f"Sentiment analysis failed: {e}")
+
         time.sleep(1)  # Задержка для предотвращения перегрузки
-    print("Анализ настроений завершен. Общие результаты: {results}")
+    print(f"Анализ настроений завершен. Общие результаты: {results}")
     return results