Update routers/textclas.py

routers/textclas.py

@@ -1,17 +1,8 @@
 from fastapi import APIRouter, Query, HTTPException
-from sentence_transformers import SentenceTransformer
-import pickle
-from sklearn.metrics.pairwise import cosine_similarity
+from keybert import KeyBERT
 
-# Carrega o modelo de embeddings
-model = SentenceTransformer('all-MiniLM-L6-v2')
-
-# Carrega os embeddings e palavras-chave pré-calculados
-with open('keywords_embeddings.pkl', 'rb') as f:
-    keywords_embeddings = pickle.load(f)
-
-with open('keywords_list.pkl', 'rb') as f:
-    keyword_categories = pickle.load(f)
+# Inicializa o modelo KeyBERT com DistilBERT (mais rápido que BERT completo)
+kw_model = KeyBERT(model='distilbert-base-nli-mean-tokens')
 
 router = APIRouter()
 
@@ -21,27 +12,31 @@ 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 em similaridade semântica.
+    Extrai palavras-chave relevantes de um texto.
     """
     try:
-        # Gera o embedding do texto
-        text_embedding = model.encode([text])
-
-        # Calcula a similaridade entre o texto e as palavras-chave
-        similarities = cosine_similarity(text_embedding, keywords_embeddings)
-
-        # Ordena as palavras-chave com base na similaridade
-        sorted_indices = similarities[0].argsort()[::-1]
-
-        # Retorna as palavras-chave com maior similaridade
-        top_keywords = [keyword_categories[i] for i in sorted_indices[:num_keywords]]
-
+        # Extrai palavras-chave
+        keywords = kw_model.extract_keywords(
+            text,
+            keyphrase_ngram_range=(1, 2),
+            stop_words='english',
+            top_n=num_keywords
+        )
+        
+        # Formata o retorno
+        keyword_list = [kw[0] for kw in keywords]
+        
         return {
             "text": text,
             "num_keywords": num_keywords,
-            "keywords": top_keywords
+            "keywords": keyword_list
         }
-
+    
+    except ValueError as ve:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid input: {str(ve)}"
+        )
     except Exception as e:
         raise HTTPException(
             status_code=500,