app.py

import gradio as gr
import torch
import torch.nn.functional as F
import os
import sys
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")
    
    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. 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()

    # --- 1. BERT Puro (Baseline) ---
    with torch.no_grad():
        outputs_base = mlm_model(**inputs)
        logits_base = outputs_base.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 = tokenizer.decode([idx]).strip()
            res_base.append(f"**{token}** ({score:.2%})")

    # --- 2. Betina (Com Vórtice) ---
    with torch.no_grad():
        # a) Gerar embedding semântico do texto todo
        emb = embedding_model.encode(texto_completo, convert_to_tensor=True).to(device)
        
        # b) Projetar para dimensão do Vórtice
        proj = embedding_projector(emb)
        
        # c) Passar pelo Vórtice (O Cérebro Caótico)
        # O método forward retorna: evolved, loss, metrics, delta_inter
        # Precisamos adicionar dimensão de batch (unsqueeze)
        # AGORA COM FATOR CAOS!
        _, _, 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 do BERT
        outputs_hidden = mlm_model(**inputs, output_hidden_states=True)
        last_hidden_state = outputs_hidden.hidden_states[-1]
        
        # Soma a correção (broadcast)
        corrected_hidden = last_hidden_state + 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 (A MIRA LASER) ---
        # Se o Caos estiver ativado, forçamos o modelo a considerar palavras do contexto.
        # Isso resolve o problema de "quebrar o viés mas não achar a resposta".
        if chaos_factor > 1.0:
            # Tokeniza apenas o contexto para descobrir quais palavras estão lá
            context_tokens = tokenizer(contexto, add_special_tokens=False)["input_ids"]
            
            # Cria um vetor de reforço
            resonance_bias = torch.zeros_like(logits_betina[0, mask_idx])
            
            # Para cada palavra que aparece no contexto, aumentamos a probabilidade dela
            for token_id in context_tokens:
                # O peso é proporcional ao Fator Caos. 
                # 0.5 é um ajuste fino para não ignorar totalmente a gramática.
                resonance_bias[token_id] += (chaos_factor * 0.5)
            
            # Injeta a ressonância nos logits originais
            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 = tokenizer.decode([idx]).strip()
            res_betina.append(f"**{token}** ({score:.2%})")

    # Formatar saída HTML (Estilo Clean)
    html_output = f"""
    <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]
        ],
        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()