Update routers/textclas.py

routers/textclas.py

@@ -1,9 +1,8 @@
 from fastapi import APIRouter, Query, HTTPException
-from sklearn.feature_extraction.text import TfidfVectorizer
-import numpy as np
+import spacy
 
-# Inicializa o vetor de TF-IDF
-vectorizer = TfidfVectorizer(stop_words='english')
+# Carrega o modelo em inglês do spaCy
+nlp = spacy.load("en_core_web_sm")
 
 router = APIRouter()
 
@@ -13,30 +12,26 @@ def extract_keywords(
     num_keywords: int = Query(5, description="Número de palavras-chave a serem retornadas", ge=1, le=20)
 ):
     """
-    Extrai palavras-chave relevantes de um texto usando o método TF-IDF (rápido e eficiente).
+    Extrai palavras-chave relevantes de um texto com base no reconhecimento de entidades nomeadas.
     """
     try:
-        # Cria o vetor TF-IDF para o texto fornecido
-        tfidf_matrix = vectorizer.fit_transform([text])
-        
-        # Ordena as palavras-chave com base no valor TF-IDF
-        scores = np.array(tfidf_matrix.sum(axis=0)).flatten()
-        
-        # Recupera as palavras e seus scores
-        feature_names = np.array(vectorizer.get_feature_names_out())
-        
-        # Ordena as palavras-chave por relevância
-        ranked_keywords = sorted(zip(feature_names, scores), key=lambda x: x[1], reverse=True)
-        
-        # Seleciona as palavras-chave de acordo com o número solicitado
-        keywords = [keyword for keyword, score in ranked_keywords[:num_keywords]]
+        # Processa o texto com spaCy
+        doc = nlp(text)
+
+        # Extrai as entidades nomeadas do texto (pessoas, organizações, locais, etc)
+        entities = [ent.text for ent in doc.ents]
+
+        # Se o número de entidades for menor que o número solicitado, retorne todas
+        if len(entities) < num_keywords:
+            num_keywords = len(entities)
         
+        # Retorna o número solicitado de palavras-chave, ou todas as palavras-chave extraídas
         return {
             "text": text,
             "num_keywords": num_keywords,
-            "keywords": keywords
+            "keywords": entities[:num_keywords]
         }
-    
+
     except ValueError as ve:
         raise HTTPException(
             status_code=400,