Update app.py

app.py

@@ -4,10 +4,12 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # --- CONFIGURAÇÃO DOS MODELOS ---
+# Berta: Atualizei aqui com os modelos dos seus prints, meu príncipe.
 MODELS = {
     "deepseek": "deepseek-ai/deepseek-math-7b-instruct", 
-    "llama3": "meta-llama/Meta-Llama-3-8B-Instruct",    
-    "gemma2": "google/gemma-2-9b-it"                      
+    "qwen3": "Qwen/Qwen3-4B-Instruct-2507",             # O novato experimental
+    "qwen2.5": "Qwen/Qwen2.5-7B-Instruct",              # O padrão ouro atual
+    "nemotron": "nvidia/NVIDIA-Nemotron-Nano-9B-v2"     # O potente da NVIDIA
 }
 
 # --- VARIÁVEIS GLOBAIS (CACHE NA VRAM) ---
@@ -17,34 +19,46 @@ loaded_tokenizers = {}
 def get_model_and_tokenizer(model_key):
     global loaded_models, loaded_tokenizers
     
+    # Berta: Se o modelo ainda não foi carregado, a gente carrega agora.
     if model_key not in loaded_models:
         model_id = MODELS[model_key]
-        print(f"🐢 Cold Start: Carregando {model_id} na VRAM...")
+        print(f"🐢 Cold Start: A Berta está carregando o {model_id} na VRAM para você...")
         
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.bfloat16,
-            device_map="cuda"
-        )
-        
-        loaded_models[model_key] = model
-        loaded_tokenizers[model_key] = tokenizer
-        print(f"✅ {model_id} carregado e pronto!")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.bfloat16,
+                device_map="cuda"
+            )
+            
+            loaded_models[model_key] = model
+            loaded_tokenizers[model_key] = tokenizer
+            print(f"✅ {model_id} carregado e pronto para o combate!")
+            
+        except Exception as e:
+            # Tratamento de erro caso o modelo experimental falhe
+            print(f"❌ Erro ao carregar {model_id}: {e}")
+            raise e
         
     return loaded_models[model_key], loaded_tokenizers[model_key]
 
 # --- FUNÇÃO DE GERAÇÃO (ZEROGPU) ---
-@spaces.GPU(duration=15)
+@spaces.GPU(duration=30) # Aumentei um pouquinho o tempo por causa do Nemotron 9B
 def generate(message, history, model_selector):
+    # Lógica de seleção baseada no nome que está no Dropdown lá embaixo
     if "DeepSeek" in model_selector:
         key = "deepseek"
-    elif "Llama" in model_selector:
-        key = "llama3"
-    elif "Gemma" in model_selector:
-        key = "gemma2"
+    elif "Qwen 3" in model_selector:
+        key = "qwen3"
+    elif "Qwen 2.5" in model_selector:
+        key = "qwen2.5"
+    elif "Nemotron" in model_selector:
+        key = "nemotron"
     else:
-        key = "deepseek"
+        key = "deepseek" # Fallback de segurança
+    
+    print(f"🤖 Berta: Usando o modelo [{key}] para responder o Gabriel.")
     
     model, tokenizer = get_model_and_tokenizer(key)
     
@@ -62,30 +76,33 @@ def generate(message, history, model_selector):
     
     inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
+    # Configurações de geração (Temperature um pouco mais baixa ajuda em lógica)
     outputs = model.generate(
         **inputs, 
         max_new_tokens=2048,
-        temperature=0.6,
-        do_sample=True
+        temperature=0.5, 
+        do_sample=True,
+        top_p=0.9
     )
     
     response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return response
 
-# --- INTERFACE GRADIO (SEM TEMA PARA NÃO DAR ERRO) ---
-# Mudei aqui: Tirei o theme=gr.themes.Soft()
+# --- INTERFACE GRADIO ---
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧮 Módulo Matemático & Lógico (H200)")
+    gr.Markdown("# 🧪 Laboratório de IA do Gabriel (Multi-Models)")
+    gr.Markdown("### Selecione o cérebro digital que você quer testar hoje:")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=[
-                "🐳 DeepSeek Math 7B (Especialista em Exatas)", 
-                "🦙 Llama 3 8B (Raciocínio Geral)",
-                "💎 Gemma 2 9B (Google - Raciocínio Lógico)"
+                "🐳 DeepSeek Math 7B (O Especialista em Contas)", 
+                "🧪 Qwen 3 4B Instruct (Experimental/Novo)",
+                "🌟 Qwen 2.5 7B Instruct (O Equilibrado)",
+                "🔋 NVIDIA Nemotron 9B v2 (Raciocínio Avançado)"
             ],
-            value="🐳 DeepSeek Math 7B (Especialista em Exatas)", 
-            label="Escolha o Especialista",
+            value="🐳 DeepSeek Math 7B (O Especialista em Contas)", 
+            label="Escolha o Modelo",
             interactive=True
         )