Update routers/textclas.py

routers/textclas.py

@@ -12,7 +12,7 @@ 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 com base no reconhecimento de entidades nomeadas.
+    Extrai palavras-chave relevantes de um texto com base em entidades nomeadas, substantivos e adjetivos.
     """
     try:
         # Processa o texto com spaCy
@@ -21,15 +21,24 @@ def extract_keywords(
         # 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)
+        # Extrai substantivos e adjetivos importantes
+        keywords = [token.lemma_ for token in doc if token.pos_ in ['NOUN', 'ADJ'] and not token.is_stop]
+
+        # Junta as entidades com as palavras-chave encontradas
+        keywords.extend(entities)
+
+        # Remover duplicatas
+        keywords = list(set(keywords))
+
+        # Se o número de palavras-chave for menor que o número solicitado, retorne todas
+        if len(keywords) < num_keywords:
+            num_keywords = len(keywords)
         
         # Retorna o número solicitado de palavras-chave, ou todas as palavras-chave extraídas
         return {
             "text": text,
             "num_keywords": num_keywords,
-            "keywords": entities[:num_keywords]
+            "keywords": keywords[:num_keywords]
         }
 
     except ValueError as ve: