Update app.py

app.py

@@ -1,68 +1,36 @@
-import spacy
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from nltk.tokenize import word_tokenize
-from nltk.corpus import stopwords
-from collections import Counter
-import re
-import os
-import nltk
-
-# Definir um diretório local para armazenar dados do NLTK
-os.environ['NLTK_DATA'] = './nltk_data'  # Dentro do diretório do projeto, por exemplo
-
-# Baixar os pacotes necessários
-nltk.download('punkt')
-nltk.download('stopwords')
-
-
-# Carregar o modelo spaCy para reconhecimento de entidades nomeadas
-nlp = spacy.load("en_core_web_sm")
+import yake
 
 app = FastAPI()
 
-class PostText(BaseModel):
-    post: str  # O post (texto) que será enviado para a API
+# Inicializando o extrator de palavras-chave YAKE
+kw_extractor = yake.KeywordExtractor()
 
-def clean_text(text):
-    """Remove caracteres especiais e faz a limpeza do texto."""
-    text = re.sub(r'[^\w\s]', '', text)  # Remove pontuação
-    text = text.lower()  # Converte para minúsculas
-    return text
-
-def extract_keywords(text):
-    """Extrai palavras-chave usando spaCy e nltk."""
-    
-    # Limpeza inicial do texto
-    cleaned_text = clean_text(text)
+@app.get("/extract-keywords")
+async def extract_keywords(text: str, max_ngram_size: int = 3, num_of_keywords: int = 10):
+    """
+    Extrai palavras-chave de um texto enviado via GET.
     
-    # Tokenização do texto
-    words = word_tokenize(cleaned_text)
+    :param text: Texto do post a ser analisado.
+    :param max_ngram_size: O tamanho máximo dos n-grams (default é 3).
+    :param num_of_keywords: Número máximo de palavras-chave a serem retornadas.
     
-    # Remover stopwords
-    stop_words = set(stopwords.words("english"))
-    filtered_words = [word for word in words if word not in stop_words]
-
-    # Contar a frequência das palavras filtradas
-    word_counts = Counter(filtered_words)
-    
-    # Processar entidades nomeadas com spaCy (ex.: pessoas, locais, eventos)
-    doc = nlp(text)
-    entities = [ent.text for ent in doc.ents]
-    
-    # Juntar as palavras mais frequentes e as entidades encontradas
-    keywords = set(filtered_words + entities)
-    
-    # Ordenar e retornar as palavras-chave mais relevantes (top 10)
-    return [keyword for keyword, _ in word_counts.most_common(10)] + entities[:10]
-
-@app.get("/generate-keywords")
-async def generate_keywords(post_text: PostText):
+    :return: Lista de palavras-chave extraídas.
+    """
     try:
-        # Gerar as palavras-chave a partir do texto do post
-        keywords = extract_keywords(post_text.post)
+        # Configuração do extrator de palavras-chave
+        custom_kw_extractor = yake.KeywordExtractor(
+            lan="en", n=max_ngram_size, top=num_of_keywords
+        )
+        
+        # Extraindo as palavras-chave
+        keywords = custom_kw_extractor.extract_keywords(text)
+        
+        # Retornando as palavras-chave (ignora os scores)
+        keyword_list = [kw[1] for kw in keywords]
         
-        return {"keywords": keywords}
+        return {"keywords": keyword_list}
     
     except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
+        raise HTTPException(status_code=400, detail=str(e))