filtro por embedding de prompt

app.py

@@ -1,17 +1,10 @@
 import gradio as gr
-from huggingface_hub import login
-import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from sentence_transformers import SentenceTransformer, util
+import torch
 
-# Login com token (via secrets do Hugging Face Spaces)
-hf_token = os.environ.get("HF_TOKEN")
-if hf_token:
-    login(hf_token)
-
-# Modelo escolhido: Google Gemma 2B Instruct
+# Modelo de linguagem (Gemma 2B IT)
 model_id = "google/gemma-2b-it"
-
-# Carregar tokenizer e modelo
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -19,53 +12,32 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype="auto"
 )
 
+# Modelo de embeddings (leve e rápido)
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+
+# Embedding do domínio fitness
+fitness_domain = "perguntas sobre treino, musculação, nutrição, hipertrofia, academia e recuperação física"
+fitness_embedding = embedder.encode(fitness_domain, convert_to_tensor=True)
+
 def responder(prompt):
-    # Lista ampliada de palavras-chave fitness
-    fitness_keywords = [
-        # treino
-        "treino", "treinos", "treinamento", "treinamentos",
-        
-        # exercício
-        "exercício", "exercicio", "exercícios", "exercicios",
-        
-        # academia
-        "academia", "musculação", "musculacao", "gym",
-        
-        # tipos de treino
-        "hipertrofia", "força", "forca", "forças", "forcas",
-        "resistência", "resistencia", "plyo", "pliométrico", "pliometrico",
-        
-        # partes do corpo
-        "perna", "pernas", "braço", "braco", "braços", "bracos",
-        "peito", "costas", "ombro", "ombros", "abdômen", "abdomen",
-        "core", "glúteo", "gluteo", "glúteos", "gluteos",
-        
-        # nutrição
-        "nutrição", "nutricao", "alimentação", "alimentacao", "dieta", "dietas",
-        "proteína", "proteina", "proteínas", "proteinas",
-        "carboidrato", "carboidratos", "gordura", "gorduras",
-        
-        # recuperação
-        "recuperação", "recuperacao", "descanso", "sono", "alongamento", "stretching",
-        
-        # cardio
-        "cardio", "corrida", "correr", "bike", "bicicleta", "ciclismo", "esteira"
-    ]
+    # Embedding do prompt
+    prompt_embedding = embedder.encode(prompt, convert_to_tensor=True)
     
-    # Checar se o prompt contém algum termo fitness
-    if not any(palavra in prompt.lower() for palavra in fitness_keywords):
+    # Similaridade com o domínio fitness
+    similarity = util.cos_sim(prompt_embedding, fitness_embedding).item()
+    
+    # Threshold ajustável (0.3–0.4 é bom para português)
+    if similarity < 0.35:
         return "Desculpe, só respondo perguntas sobre treino, nutrição e fitness."
-
+    
     # Prompt explícito
     system_message = (
         "Você é um personal trainer virtual. "
         "Sempre responda em PORTUGUÊS, de forma clara, curta e prática. "
         "Se o usuário pedir treino, forneça uma lista numerada de exercícios físicos reais. "
-        "Se não for relacionado a treino/nutrição/fitness, responda 'Desculpe, só respondo perguntas sobre treino e fitness.'"
     )
 
     entrada = f"{system_message}\n\nUsuário: {prompt}\nAssistente:"
-
     inputs = tokenizer(entrada, return_tensors="pt").to(model.device)
 
     outputs = model.generate(
@@ -86,8 +58,7 @@ demo = gr.Interface(
     fn=responder,
     inputs=gr.Textbox(lines=3, label="Pergunta"),
     outputs=gr.Textbox(label="Resposta"),
-    title="Personal Trainer AI"
+    title="Personal Trainer AI com Filtro Semântico"
 )
 
-# Ativa fila (compatível com versões antigas do Gradio no Spaces)
 demo.queue().launch()