Update clustering.py

clustering.py

@@ -28,8 +28,15 @@ from wordcloud import WordCloud
 
 from concurrent.futures import ThreadPoolExecutor
 
-if os.getenv("RAILWAY_ENVIRONMENT") is None:
-    load_dotenv()
+try:
+    import torch
+    device = 0 if torch.cuda.is_available() else -1
+    embeddings_device = "cuda"
+    batch_size = 128
+except ImportError:
+    device = -1  # Si no está instalado, forzar uso de CPU
+    embeddings_device = "cpu"
+    batch_size = 32
 
 api_key = os.getenv("youtube_api_key")
 
@@ -338,18 +345,18 @@ def add_normalized_embeddings_to_dataframe(
     Parámetros:
     - data: pandas DataFrame, el DataFrame que contiene la columna de texto.
     - text_column: str, el nombre de la columna en el DataFrame que contiene el texto para generar embeddings.
-    - model_name: str, el nombre del modelo de SentenceTransformer a utilizar (por defecto "sentence-transformers/stsb-xlm-r-multilingual").
+    - model_name: str, el nombre del modelo de SentenceTransformer a utilizar.
+    - batch_size: int, el tamaño del lote para procesamiento eficiente.
 
     Retorna:
     - data: pandas DataFrame, el DataFrame original con las nuevas columnas de embeddings normalizados.
     """
 
-    model = SentenceTransformer(model_name)
+    model = SentenceTransformer(model_name, device=embeddings_device)
     sentences = data[text_column].tolist()
-    embeddings = model.encode(sentences)
+    embeddings = model.encode(sentences, batch_size=batch_size, convert_to_numpy=True, device=device)
     normalized_embeddings = normalize(embeddings, norm="l2")
-
-    data["embeddings"] = [embedding for embedding in normalized_embeddings]
+    data["embeddings"] = list(normalized_embeddings)
 
     return data
 
@@ -854,6 +861,7 @@ classifier = pipeline(
     "sentiment-analysis",
     model="nlptown/bert-base-multilingual-uncased-sentiment",
     truncation=True,
+    device=device
 )
 
 
@@ -866,22 +874,25 @@ def map_sentiment(estrella):
         return "positivo"
 
 
-def classify_sentiment_df(data, comment_col="comment", num_threads=8):
-    def classify_sentiment(texto):
-        resultado = classifier(texto)[0]
-        sentimiento = map_sentiment(resultado["label"])
-        return sentimiento, resultado["score"]
+def classify_sentiment_df(data, comment_col="comment", batch_size=32, num_threads=8):
+    comentarios = data[comment_col].tolist()
 
-    def process_row(texto):
-        # Esta función procesa una fila y retorna los valores necesarios
-        return classify_sentiment(texto)
+    if device == 0:  # Si hay GPU, procesar en batch
+        resultados = classifier(comentarios, batch_size=batch_size)
+        data["sentimiento"] = [map_sentiment(r["label"]) for r in resultados]
+        data["confianza"] = [r["score"] for r in resultados]
+    
+    else:  # Si no hay GPU, usar CPU con hilos
+        def classify_sentiment(texto):
+            resultado = classifier(texto)[0]
+            return map_sentiment(resultado["label"]), resultado["score"]
 
-    # Crear un ThreadPoolExecutor con el número de hilos especificado
-    with ThreadPoolExecutor(max_workers=num_threads) as executor:
-        results = list(executor.map(process_row, data[comment_col]))
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            resultados = list(executor.map(classify_sentiment, comentarios))
 
-    # Dividir los resultados en columnas para "sentimiento" y "confianza"
-    data["sentimiento"], data["confianza"] = zip(*results)
+        sentimientos, confianzas = zip(*resultados)
+        data["sentimiento"] = sentimientos
+        data["confianza"] = confianzas
 
     return data