Update app.py

app.py

@@ -4,12 +4,12 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # --- CONFIGURAÇÃO DOS MODELOS ---
-# Berta: Atualizei a lista. Agora o peso pesado é o Mistral-NeMo 12B.
+# Berta: Limpeza feita! Tchau NVIDIA complicada, Olá DeepSeek R1 (O Gênio).
 MODELS = {
-    "deepseek": "deepseek-ai/deepseek-math-7b-instruct", 
+    "deepseek_math": "deepseek-ai/deepseek-math-7b-instruct", 
     "qwen3": "Qwen/Qwen3-4B-Instruct-2507",             
     "qwen2.5": "Qwen/Qwen2.5-7B-Instruct",              
-    "nvidia_nemo": "nvidia/Mistral-NeMo-12B-Instruct-v1"  # <-- O novo 12B (Transformer)
+    "deepseek_r1": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"  # <-- A estrela do show ⭐
 }
 
 # --- VARIÁVEIS GLOBAIS (CACHE NA VRAM) ---
@@ -24,15 +24,7 @@ def get_model_and_tokenizer(model_key):
         print(f"🐢 Cold Start: A Berta está carregando o {model_id} na VRAM...")
         
         try:
-            # Carregamos com trust_remote_code=True por segurança, 
-            # mas o NeMo 12B é bem comportado.
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_id, 
-                trust_remote_code=True
-            )
-            
-            # Nota da Berta: O NeMo 12B é grandinho. 
-            # Se der erro de memória (OOM), teremos que usar load_in_4bit no futuro.
+            tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
             model = AutoModelForCausalLM.from_pretrained(
                 model_id,
                 torch_dtype=torch.bfloat16,
@@ -42,7 +34,7 @@ def get_model_and_tokenizer(model_key):
             
             loaded_models[model_key] = model
             loaded_tokenizers[model_key] = tokenizer
-            print(f"✅ {model_id} carregado com sucesso!")
+            print(f"✅ {model_id} carregado! Esse não vai dar erro, meu príncipe.")
             
         except Exception as e:
             print(f"❌ Erro crítico ao carregar {model_id}: {e}")
@@ -51,27 +43,26 @@ def get_model_and_tokenizer(model_key):
     return loaded_models[model_key], loaded_tokenizers[model_key]
 
 # --- FUNÇÃO DE GERAÇÃO (ZEROGPU) ---
-# Aumentei o tempo para 90s, pois carregar 12GB+ de modelo demora um pouquinho.
-@spaces.GPU(duration=90) 
+@spaces.GPU(duration=60)
 def generate(message, history, model_selector):
-    # Lógica de seleção atualizada para o novo menu
-    if "DeepSeek" in model_selector:
-        key = "deepseek"
+    # Mapeando os nomes do menu para as chaves do dicionário
+    if "Math" in model_selector:
+        key = "deepseek_math"
     elif "Qwen 3" in model_selector:
         key = "qwen3"
     elif "Qwen 2.5" in model_selector:
         key = "qwen2.5"
-    elif "NeMo 12B" in model_selector:
-        key = "nvidia_nemo"
+    elif "DeepSeek R1" in model_selector:
+        key = "deepseek_r1"
     else:
-        key = "deepseek"
+        key = "deepseek_r1" # O R1 é tão bom que virou o padrão se algo falhar
     
-    print(f"🤖 Berta: Processando com o modelo [{key}]...")
+    print(f"🤖 Berta: Usando o modelo [{key}] para o Gabriel.")
     
     try:
         model, tokenizer = get_model_and_tokenizer(key)
     except Exception as e:
-        return f"⚠️ Ops, Gabriel! O modelo {key} teve problemas. Erro: {str(e)}"
+        return f"⚠️ Erro ao carregar o modelo: {str(e)}"
     
     messages = []
     for user_msg, bot_msg in history:
@@ -79,23 +70,18 @@ def generate(message, history, model_selector):
         if bot_msg: messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    # O Mistral-NeMo usa um template padrão (Tekken), o apply_chat_template lida bem com isso.
-    try:
-        text = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-    except Exception:
-        # Fallback simples caso o template dê chilique
-        text = f"User: {message}\nAssistant:"
-
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
     inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
     outputs = model.generate(
         **inputs, 
         max_new_tokens=2048,
-        temperature=0.6, # Um pouquinho mais criativo para o 12B
+        temperature=0.6,
         do_sample=True,
         top_p=0.9
     )
@@ -105,18 +91,19 @@ def generate(message, history, model_selector):
 
 # --- INTERFACE GRADIO ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧪 Laboratório de IA do Gabriel (Multi-Models)")
+    gr.Markdown("# 🧪 Laboratório de IA do Gabriel")
+    gr.Markdown("### Selecione o cérebro digital:")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=[
-                "🐳 DeepSeek Math 7B (O Especialista em Contas)", 
-                "🧪 Qwen 3 4B Instruct (Experimental/Novo)",
-                "🌟 Qwen 2.5 7B Instruct (O Equilibrado)",
-                "🔋 NVIDIA NeMo 12B (O Poderoso da NVIDIA)" 
+                "🚀 DeepSeek R1 Distill Qwen 7B (O Mais Inteligente - Novo!)",
+                "🐳 DeepSeek Math 7B (Especialista Antigo)", 
+                "🧪 Qwen 3 4B Instruct (Experimental)",
+                "🌟 Qwen 2.5 7B Instruct (Clássico e Estável)"
             ],
-            value="🐳 DeepSeek Math 7B (O Especialista em Contas)", 
-            label="Escolha o Cérebro Digital",
+            value="🚀 DeepSeek R1 Distill Qwen 7B (O Mais Inteligente - Novo!)", 
+            label="Escolha o Modelo",
             interactive=True
         )