Uniaff commited on
Commit
8471142
·
verified ·
1 Parent(s): f2d0404

Update func_ai.py

Browse files
Files changed (1) hide show
  1. func_ai.py +11 -8
func_ai.py CHANGED
@@ -1,19 +1,17 @@
1
-
2
  import requests
3
  import torch
4
- # from googletrans import Translator
5
  from transformers import pipeline
6
  from deep_translator import GoogleTranslator
7
  import time
8
  import os
9
- VECTOR_API_URL = os.getenv('API_URL')
10
 
11
- # translator = Translator()
12
 
 
13
  sentiment_model = pipeline(
14
  'sentiment-analysis',
15
- model='cardiffnlp/twitter-xlm-roberta-base-sentiment',
16
- tokenizer='cardiffnlp/twitter-xlm-roberta-base-sentiment',
17
  device=0 if torch.cuda.is_available() else -1
18
  )
19
 
@@ -66,9 +64,14 @@ def analyze_sentiment(comments):
66
  for i in range(0, len(comments), 50):
67
  batch = comments[i:i + 50]
68
  print(f"Анализируем батч с {i} по {i + len(batch)} комментарий: {batch}")
69
- batch_results = sentiment_model(batch)
 
 
 
 
 
70
  print(f"Результаты батча: {batch_results}")
71
  results.extend(batch_results)
72
  time.sleep(1) # Задержка для предотвращения перегрузки
73
- print("Анализ настроений завершен. Общие результаты: {results}")
74
  return results
 
 
1
  import requests
2
  import torch
 
3
  from transformers import pipeline
4
  from deep_translator import GoogleTranslator
5
  import time
6
  import os
 
7
 
8
+ VECTOR_API_URL = os.getenv('API_URL')
9
 
10
+ # Replace the sentiment analysis model
11
  sentiment_model = pipeline(
12
  'sentiment-analysis',
13
+ model='distilbert-base-uncased-finetuned-sst-2-english',
14
+ tokenizer='distilbert-base-uncased-finetuned-sst-2-english',
15
  device=0 if torch.cuda.is_available() else -1
16
  )
17
 
 
64
  for i in range(0, len(comments), 50):
65
  batch = comments[i:i + 50]
66
  print(f"Анализируем батч с {i} по {i + len(batch)} комментарий: {batch}")
67
+
68
+ # Translate comments into English before sentiment analysis
69
+ translated_batch = [GoogleTranslator(source='auto', target="en").translate(comment) for comment in batch]
70
+ print(f"Переведённый батч: {translated_batch}")
71
+
72
+ batch_results = sentiment_model(translated_batch)
73
  print(f"Результаты батча: {batch_results}")
74
  results.extend(batch_results)
75
  time.sleep(1) # Задержка для предотвращения перегрузки
76
+ print(f"Анализ настроений завершен. Общие результаты: {results}")
77
  return results