new_ads_obnova / func_ai.py
Uniaff's picture
Update func_ai.py
ea6fb88 verified
raw
history blame
1.96 kB
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()
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
)
def classify_comment(text):
translated_text = GoogleTranslator(source='auto', target=target_language).translate(text)
result = classifier(translated_text, ["interrogative", "non-interrogative"], clean_up_tokenization_spaces=True)
top_class = result['labels'][0]
return top_class
def retrieve_from_vdb(query):
print(f"Отправка запроса к FastAPI сервису: {query}")
response = requests.get(f"{VECTOR_API_URL}/search/", json={"query": query})
if response.status_code == 200:
results = response.json().get("results", [])
print(f"Получено {len(results)} результатов.")
return results
else:
print(f"Ошибка при поиске: {response.text}")
return []
def analyze_sentiment(comments):
print("Начинаем анализ настроений.")
results = []
for i in range(0, len(comments), 50):
batch = comments[i:i + 50]
print(f"Анализируем батч с {i} по {i + len(batch)} комментарий.")
batch_results = sentiment_model(batch)
results.extend(batch_results)
time.sleep(1) # Задержка для предотвращения перегрузки
print("Анализ настроений завершен.")
return results