app.py

import gradio as gr
import torch
import torch.nn.functional as F
import os
import sys
import re
from functools import lru_cache
from transformers import AutoModelForMaskedLM, AutoTokenizer
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download

# Importa a classe real do seu arquivo bettina.py
# Certifique-se de que bettina.py está na mesma pasta
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
try:
    from bettina import VortexBetinaAntiHalluc
except ImportError:
    # Tenta importar assumindo que estamos na raiz do projeto
    try:
        import bettina
        VortexBetinaAntiHalluc = bettina.VortexBetinaAntiHalluc
    except ImportError as e:
        raise ImportError(f"CRÍTICO: Não foi possível encontrar 'bettina.py'. Verifique se o arquivo foi enviado para o Space. Erro: {e}")

# Configuração de Dispositivo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Rodando em: {device}")

# ==============================================================================
# 1. Carregamento dos Modelos Base
# ==============================================================================
print("Carregando modelos base...")
embedding_model_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
tokenizer_name = "neuralmind/bert-base-portuguese-cased"

# Carrega modelos com cache para não baixar toda vez
embedding_model = SentenceTransformer(embedding_model_name, device=str(device))
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
mlm_model = AutoModelForMaskedLM.from_pretrained(tokenizer_name).to(device)
mlm_model.eval()

# ==============================================================================
# 2. Inicialização da Betina (Nosso Cérebro)
# ==============================================================================
# Configurações devem bater com o que foi treinado. Usando defaults do bettina.py
EMBED_DIM = 256
RAW_EMBED_DIM = embedding_model.get_sentence_embedding_dimension() # 768
HIDDEN_SIZE = mlm_model.config.hidden_size # 768

print("Inicializando Vortex Betina...")
# Instancia a classe robusta do seu código
vortex = VortexBetinaAntiHalluc(
    embed_dim=EMBED_DIM,
    # Habilitando recursos avançados por padrão para demonstração
    enable_rotation=True,
    enable_quadratic_reflection=True,
    enable_lorentz_transform=True,
    enforce_square_geometry=True
).to(device)

# Projetores para conectar os mundos (SentenceTransformer -> Vortex -> BERT)
embedding_projector = torch.nn.Linear(RAW_EMBED_DIM, EMBED_DIM).to(device)
correction_projector = torch.nn.Linear(EMBED_DIM, HIDDEN_SIZE).to(device)

# ==============================================================================
# 3. Carregamento de Pesos (Se existirem)
# ==============================================================================
weights_loaded = False
REPO_ID = "reynaldo22/betina-perfect-2025"

# 1. Tentar baixar do Hugging Face Hub
try:
    print(f"Tentando baixar pesos do repositório: {REPO_ID}...")
    token = os.getenv("HF_TOKEN")
    
    # Fallback para arquivo local (para quem não consegue criar ENV)
    if not token and os.path.exists("token.txt"):
        try:
            with open("token.txt", "r") as f:
                token = f.read().strip()
            print("⚠️ Usando token do arquivo token.txt")
        except Exception as token_file_error:
            print(f"⚠️ Falha ao ler token.txt: {token_file_error}")

    vortex_path = hf_hub_download(repo_id=REPO_ID, filename="vortex.pt", token=token)
    emb_path = hf_hub_download(repo_id=REPO_ID, filename="embedding_projector.pt", token=token)
    corr_path = hf_hub_download(repo_id=REPO_ID, filename="correction_projector.pt", token=token)
    
    # strict=False permite carregar pesos parciais se houver pequenas diferenças de versão
    vortex.load_state_dict(torch.load(vortex_path, map_location=device), strict=False)
    embedding_projector.load_state_dict(torch.load(emb_path, map_location=device))
    correction_projector.load_state_dict(torch.load(corr_path, map_location=device))
    weights_loaded = True
    print("✅ Pesos carregados do Hugging Face com sucesso!")
except Exception as e:
    print(f"⚠️ Falha ao baixar do Hugging Face: {e}")
    print("Tentando carregar localmente...")

# 2. Fallback para arquivos locais
if not weights_loaded:
    POSSIBLE_DIRS = ["outputs/betina_vortex", ".", "model_weights"]
    for model_dir in POSSIBLE_DIRS:
        vortex_path = os.path.join(model_dir, "vortex.pt")
        if os.path.exists(vortex_path):
            print(f"Carregando pesos locais de {model_dir}...")
            try:
                vortex.load_state_dict(torch.load(vortex_path, map_location=device))
                embedding_projector.load_state_dict(torch.load(os.path.join(model_dir, "embedding_projector.pt"), map_location=device))
                correction_projector.load_state_dict(torch.load(os.path.join(model_dir, "correction_projector.pt"), map_location=device))
                weights_loaded = True
                break
            except Exception as e:
                print(f"Erro ao carregar pesos de {model_dir}: {e}")

if not weights_loaded:
    print("⚠️ AVISO: Pesos treinados não encontrados. Usando inicialização aleatória.")
    print("O modelo vai rodar, mas as respostas da Betina serão aleatórias até você treinar.")

vortex.eval()
embedding_projector.eval()
correction_projector.eval()

# ==============================================================================
# 4. Funções de Cache (Otimização de Performance)
# ==============================================================================
@lru_cache(maxsize=128)
def _cached_embedding(texto: str):
    """Cache de embeddings semânticos para evitar recomputação."""
    with torch.no_grad():
        emb = embedding_model.encode(texto, convert_to_tensor=True).to(device)
    return emb

@lru_cache(maxsize=256)
def _cached_tokenize(texto: str):
    """Cache de tokenização para textos repetidos."""
    return tuple(tokenizer(texto, add_special_tokens=False)["input_ids"])


def _pretty_token(token_id, contexto):
    """Remove marcas de subword e tenta casar fragmentos com palavras do contexto."""
    token_piece = tokenizer.convert_ids_to_tokens(int(token_id))
    
    # Tokens especiais ou vazios
    if not token_piece or token_piece in ["[UNK]", "[PAD]", "[CLS]", "[SEP]"]:
        return tokenizer.decode([int(token_id)]).strip() or "?"
    
    # Remove marcador de subword
    clean_piece = token_piece.replace("##", "").strip()
    if not clean_piece:
        return tokenizer.decode([int(token_id)]).strip() or "?"

    # Tenta casar com palavras do contexto (prioriza match exato)
    token_lower = clean_piece.lower()
    context_words = re.findall(r"\w+", contexto)
    
    # 1. Match exato primeiro
    for word in context_words:
        if word.lower() == token_lower:
            return word
    
    # 2. Match por prefixo (palavra começa com o token)
    for word in context_words:
        if word.lower().startswith(token_lower) and len(word) - len(token_lower) <= 4:
            return word
    
    # 3. Match por sufixo (token é final de palavra)
    for word in context_words:
        if word.lower().endswith(token_lower) and len(word) - len(token_lower) <= 4:
            return word

    return clean_piece

# ==============================================================================
# 4. Lógica de Inferência
# ==============================================================================
def predict(contexto, frase_mask, chaos_factor):
    if "[MASK]" not in frase_mask:
        return "⚠️ Erro: A frase precisa conter o token [MASK]."

    # Combinar contexto e frase para o embedding semântico
    texto_completo = f"{contexto} {frase_mask}".strip()
    
    # Preparar inputs para o BERT
    inputs = tokenizer(texto_completo, return_tensors="pt").to(device)
    
    # Encontrar índice da máscara
    mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
    if len(mask_token_index) == 0:
         return "Erro: Token [MASK] não identificado corretamente pelo tokenizer."
    mask_idx = mask_token_index[0].item()

    # --- Inferência Única (BERT + Betina) ---
    with torch.no_grad():
        outputs = mlm_model(**inputs, output_hidden_states=True, return_dict=True)
        logits_base = outputs.logits
        probs_base = F.softmax(logits_base[0, mask_idx], dim=-1)
        top_k_base = torch.topk(probs_base, 5)
        
        res_base = []
        for idx, score in zip(top_k_base.indices, top_k_base.values):
            token = _pretty_token(idx.item(), contexto)
            res_base.append(f"**{token}** ({score:.2%})")

        # a) Gerar embedding semântico do texto todo (COM CACHE)
        emb = _cached_embedding(texto_completo)
        
        # b) Projetar para dimensão do Vórtice
        proj = embedding_projector(emb)
        
        # c) Passar pelo Vórtice (O Cérebro Caótico)
        _, _, metrics, delta = vortex(proj.unsqueeze(0), chaos_factor=chaos_factor) 
        
        # d) Projetar correção de volta para dimensão do BERT
        correction = correction_projector(delta).unsqueeze(1) # [1, 1, hidden_size]
        
        # e) Injetar nos hidden states apenas na posição mascarada
        last_hidden_state = outputs.hidden_states[-1]
        corrected_hidden = last_hidden_state.clone()
        corrected_hidden[:, mask_idx:mask_idx+1, :] += correction
        
        # f) Predição final
        if hasattr(mlm_model, "cls"):
            logits_betina = mlm_model.cls(corrected_hidden)
        else:
            logits_betina = mlm_model.get_output_embeddings()(corrected_hidden)

        # --- 🚀 RESSONÂNCIA CONTEXTUAL + SEMÂNTICA ---
        # Fase 1: Boost para palavras LITERAIS do contexto
        # Fase 2: Boost para palavras SEMANTICAMENTE RELACIONADAS
        if chaos_factor > 1.0:
            # Tokeniza apenas o contexto para descobrir quais palavras estão lá (COM CACHE)
            context_tokens = list(_cached_tokenize(contexto))
            
            # Cria um vetor de reforço normalizado
            resonance_bias = torch.zeros_like(logits_betina[0, mask_idx])
            filtered_tokens = []
            for token_id in context_tokens:
                token_str = tokenizer.convert_ids_to_tokens(token_id)
                if token_str.startswith("##") or len(token_str) < 2:
                    continue
                filtered_tokens.append(token_id)

            unique_tokens = list(set(filtered_tokens))
            if unique_tokens:
                boost_value = (chaos_factor * 0.3) / len(unique_tokens)  # Reduzido para dar espaço ao semântico
                boost_value = min(boost_value, 3.0)
                for token_id in unique_tokens:
                    resonance_bias[token_id] += boost_value
            
            # --- 🧠 FASE 2: RESSONÂNCIA SEMÂNTICA ---
            # Encontra tokens semanticamente relacionados ao contexto
            # usando similaridade de embeddings
            context_emb = _cached_embedding(contexto)
            
            # Palavras-chave para buscar relações (extraídas do contexto)
            context_words = list(set(re.findall(r"\b[a-záéíóúàâêôãõç]{4,}\b", contexto.lower())))
            
            # Para cada palavra do contexto, encontra tokens relacionados
            semantic_candidates = []
            for word in context_words[:5]:  # Limita a 5 palavras principais
                word_emb = _cached_embedding(word)
                # Calcula similaridade com embedding do contexto completo
                sim = F.cosine_similarity(word_emb.unsqueeze(0), context_emb.unsqueeze(0)).item()
                if sim > 0.3:  # Palavra relevante
                    semantic_candidates.append(word)
            
            # Adiciona palavras relacionadas ao campo semântico
            # Mapeia conceitos comuns (felicidade->peso, calor->frio, etc)
            semantic_expansions = {
                # Estados emocionais ↔ físicos
                "felicidade": ["pesado", "leve", "gordo", "magro", "cheio", "vazio", "quilos", "peso", "grande", "pequeno"],
                "feliz": ["pesado", "leve", "gordo", "magro", "cheio", "vazio", "grande", "pequeno"],
                "triste": ["leve", "pesado", "magro", "vazio", "pequeno"],
                "quilos": ["pesado", "leve", "gordo", "magro", "peso", "massa", "grande", "pequeno"],
                "medida": ["pesado", "leve", "grande", "pequeno", "alto", "baixo", "largo", "estreito"],
                "peso": ["pesado", "leve", "gordo", "magro", "quilos", "gramas"],
                
                # Temperatura e clima
                "calor": ["quente", "frio", "gelado", "fervendo", "morno", "aquecido", "congelado"],
                "frio": ["gelado", "quente", "congelado", "aquecido", "fervendo", "morno"],
                "quente": ["frio", "gelado", "fervendo", "morno", "aquecido"],
                "sol": ["calor", "frio", "luz", "escuro", "quente", "gelado", "congelado"],
                "noite": ["escuro", "claro", "frio", "quente", "dia", "calor"],
                "congela": ["frio", "gelado", "congelado", "quente", "calor"],
                
                # Vida e morte  
                "vida": ["morte", "morrer", "nascer", "viver", "morto", "vivo", "começo", "fim"],
                "morte": ["vida", "viver", "nascer", "morrer", "morto", "vivo", "fim", "começo"],
                "morto": ["vivo", "vida", "morte", "nascer"],
                "vivo": ["morto", "morte", "vida", "morrer"],
                "nascer": ["morrer", "viver", "morte", "vida"],
                
                # Lógica e verdade
                "mentira": ["verdade", "falso", "certo", "errado", "real", "honesto", "enganar"],
                "verdade": ["mentira", "falso", "certo", "real", "honesto", "verdadeiro"],
                "falso": ["verdadeiro", "certo", "real", "mentira"],
                "pergunta": ["resposta", "responder", "questão", "dizer"],
                
                # Ciclos e reversões
                "reverso": ["inverso", "contrário", "oposto", "normal"],
                "ciclo": ["começo", "fim", "início", "término", "volta"],
                
                # 💰 Economia e dinheiro
                "gastar": ["poupar", "economizar", "guardar", "investir", "perder", "ganhar"],
                "poupar": ["gastar", "desperdiçar", "usar", "consumir", "perder"],
                "rico": ["pobre", "milionário", "falido", "endividado", "próspero", "gastar", "perder"],
                "pobre": ["rico", "milionário", "próspero", "abastado", "guardar", "poupar"],
                "dinheiro": ["gastar", "poupar", "investir", "perder", "ganhar", "economizar"],
                "economia": ["gastar", "poupar", "investir", "lucro", "prejuízo"],
                "invertida": ["normal", "contrário", "oposto", "reverso"],
                
                # 🔄 Palavras de inversão contextual (mundo ao contrário)
                "contrário": ["inverso", "oposto", "reverso", "gastar", "perder", "falhar"],
                "inverso": ["contrário", "oposto", "normal", "gastar", "perder"],
                "bizarro": ["estranho", "invertido", "contrário", "oposto"],
                "estranho": ["bizarro", "invertido", "contrário", "oposto"],
                
                # 🏥 Saúde invertida
                "saudável": ["doente", "enfermo", "mal", "pior", "adoecer"],
                "doente": ["saudável", "curado", "bem", "melhor", "sarar"],
                "adoecem": ["curam", "sarar", "curar", "melhorar", "doente", "enfermo"],
                "curam": ["adoecem", "pioram", "doente", "enfermo"],
                "tratamento": ["doente", "enfermo", "curar", "piorar", "adoecer"],
                
                # 📚 Educação invertida
                "reprovar": ["passar", "aprovar", "sucesso", "acertar"],
                "aprovar": ["reprovar", "falhar", "errar", "fracassar"],
                "inteligentes": ["tolos", "burros", "ignorantes"],
                "tolos": ["inteligentes", "gênios", "sábios"],
            }
            
            # 🔥 DETECÇÃO DE MUNDO INVERTIDO
            # Se detectar palavras de inversão, ativa boost agressivo nos opostos
            inversion_markers = {"contrário", "inverso", "invertido", "invertida", "bizarro", "estranho", "avesso", "oposto"}
            context_lower = contexto.lower()
            is_inverted_world = any(marker in context_lower for marker in inversion_markers)
            
            # Mapeamento de inversões diretas (quando em mundo invertido)
            if is_inverted_world:
                inversion_map = {
                    "rico": ["gastar", "perder", "desperdiçar", "jogar"],
                    "pobre": ["guardar", "poupar", "economizar", "investir"],
                    "saudável": ["doente", "enfermo", "mal", "pior"],
                    "doente": ["curado", "saudável", "bem", "melhor"],
                    "aprovar": ["reprovar", "falhar", "errar"],
                    "reprovar": ["passar", "aprovar", "acertar"],
                    "melhor": ["pior", "doente", "mal"],
                    "pior": ["melhor", "curado", "bem"],
                }
                # Adiciona inversões ao expanded_words com boost extra
                for word in context_words:
                    word_lower = word.lower()
                    if word_lower in inversion_map:
                        for inv_word in inversion_map[word_lower]:
                            inv_tokens = tokenizer(inv_word, add_special_tokens=False)["input_ids"]
                            for tok_id in inv_tokens:
                                tok_str = tokenizer.convert_ids_to_tokens(tok_id)
                                if not tok_str.startswith("##") and len(tok_str) >= 2:
                                    resonance_bias[tok_id] += chaos_factor * 0.8  # Boost forte!
            
            # Aplica expansão semântica
            expanded_words = set()
            for word in context_words:
                word_lower = word.lower()
                if word_lower in semantic_expansions:
                    expanded_words.update(semantic_expansions[word_lower])
            
            # Tokeniza e dá boost nas palavras expandidas
            if expanded_words:
                semantic_boost = (chaos_factor * 0.4) / max(len(expanded_words), 1)
                semantic_boost = min(semantic_boost, 4.0)
                for exp_word in expanded_words:
                    exp_tokens = tokenizer(exp_word, add_special_tokens=False)["input_ids"]
                    for tok_id in exp_tokens:
                        tok_str = tokenizer.convert_ids_to_tokens(tok_id)
                        if not tok_str.startswith("##") and len(tok_str) >= 2:
                            resonance_bias[tok_id] += semantic_boost
            
            logits_betina[0, mask_idx] += resonance_bias

        probs_betina = F.softmax(logits_betina[0, mask_idx], dim=-1)
        top_k_betina = torch.topk(probs_betina, 5)
        
        res_betina = []
        for idx, score in zip(top_k_betina.indices, top_k_betina.values):
            token = _pretty_token(idx.item(), contexto)
            res_betina.append(f"**{token}** ({score:.2%})")

    # Calcular divergência entre as respostas
    top_bert = tokenizer.decode([top_k_base.indices[0].item()]).strip()
    top_betina = tokenizer.decode([top_k_betina.indices[0].item()]).strip()
    divergiu = top_bert.lower() != top_betina.lower()
    
    divergencia_html = ""
    if divergiu:
        divergencia_html = f"""

        <div style="background: linear-gradient(90deg, #52c41a 0%, #1890ff 100%); color: white; padding: 10px; border-radius: 8px; margin-bottom: 15px; text-align: center;">

            <strong>✨ DIVERGÊNCIA DETECTADA!</strong> BERT → "{top_bert}" | Betina → "{top_betina}"

        </div>

        """
    else:
        divergencia_html = f"""

        <div style="background: #faad14; color: white; padding: 10px; border-radius: 8px; margin-bottom: 15px; text-align: center;">

            <strong>⚡ Mesma resposta principal:</strong> "{top_bert}" (aumente o Caos para forçar divergência)

        </div>

        """

    # Formatar saída HTML (Estilo Clean)
    html_output = f"""

    {divergencia_html}

    <div style="display: flex; gap: 20px; flex-wrap: wrap;">

        <div style="flex: 1; min-width: 300px; background-color: #f5f5f5; padding: 15px; border-radius: 10px; border: 1px solid #ddd;">

            <h3 style="color: #555; margin-top: 0;">🧠 BERT Padrão</h3>

            <p style="font-size: 0.9em; color: #666;"><i>O que o modelo "decorou" do treino original.</i></p>

            <ol>

                {''.join([f'<li>{item}</li>' for item in res_base])}

            </ol>

        </div>

        <div style="flex: 1; min-width: 300px; background-color: #e6f7ff; padding: 15px; border-radius: 10px; border: 2px solid #1890ff;">

            <h3 style="color: #0050b3; margin-top: 0;">🌀 Betina 2.0</h3>

            <p style="font-size: 0.9em; color: #0050b3;"><i>Correção Dinâmica (Caos: {chaos_factor}x)</i></p>

            <ol>

                {''.join([f'<li>{item}</li>' for item in res_betina])}

            </ol>

        </div>

    </div>

    <br>

    <details>

        <summary style="cursor: pointer; color: #888;">📊 Métricas do Vórtice (Estado Interno)</summary>

        <pre style="font-size: 0.8em; background: #333; color: #0f0; padding: 10px; border-radius: 5px; overflow-x: auto;">{str(metrics)}</pre>

    </details>

    """
    return html_output

# ==============================================================================
# 5. Interface Gradio
# ==============================================================================
custom_css = """

footer {visibility: hidden}

"""

with gr.Blocks(title="Betina 2.0 - Protocolo Impossível") as demo:
    gr.Markdown("""

    # 🌀 BETINA 2.0: PROTOCOLO IMPOSSÍVEL

    

    Sistema de correção neural baseado em **Dinâmica de Vórtice**.

    Aumente o **Fator Caos** para forçar a lógica sobre a estatística.

    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            txt_contexto = gr.Textbox(
                label="1. CONTEXTO (A Verdade Absoluta)", 
                placeholder="Ex: A felicidade é medida em quilos. Se estou feliz, estou...", 
                lines=3
            )
            txt_mask = gr.Textbox(
                label="2. CONSULTA (Use [MASK])", 
                placeholder="Ex: Estou muito feliz, logo estou [MASK].", 
                lines=2
            )
            slider_chaos = gr.Slider(
                minimum=1.0, 
                maximum=50.0, 
                value=1.0, 
                step=0.5, 
                label="🔥 FATOR CAOS (Overdrive)",
                info="1.0 = Padrão. Aumente para forçar correções impossíveis."
            )
            btn_run = gr.Button("🌀 INICIAR VÓRTICE", variant="primary")
            
        with gr.Column(scale=1):
            out_result = gr.HTML(label="Resultado Comparativo")
    
    gr.Markdown("### 🧪 Testes de Paradoxo")
    gr.Examples(
        examples=[
            ["A felicidade é medida em quilos. Se estou feliz, estou...", "Estou muito feliz, logo estou [MASK].", 10.0],
            ["Neste mundo, o gelo é quente e o fogo é frio.", "Toquei no fogo e senti [MASK].", 15.0],
            ["O ciclo da vida é reverso: morremos, vivemos e nascemos.", "Depois de viver muito, eu vou [MASK].", 20.0],
            ["Neste planeta a noite traz calor extremo e o sol congela tudo.", "Quando o sol nasce, as pessoas sentem [MASK].", 25.0],
            ["Neste laboratório, toda pergunta precisa ser respondida com a mentira exata que a transforma em verdade. Só a mentira perfeita libera o corredor.", "Quando o cientista ouve a pergunta final, ele responde com a mentira que [MASK].", 30.0],
        ],
        inputs=[txt_contexto, txt_mask, slider_chaos]
    )

    btn_run.click(fn=predict, inputs=[txt_contexto, txt_mask, slider_chaos], outputs=out_result)


if __name__ == "__main__":
    demo.launch()