Create app.py

app.py

@@ -0,0 +1,106 @@
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# --- CONFIGURAÇÃO DOS MODELOS ---
+# IDs Oficiais do Hugging Face
+MODELS = {
+    "deepseek": "deepseek-ai/deepseek-math-7b-instruct",
+    "llama3": "meta-llama/Meta-Llama-3-8B-Instruct",
+    "gemma2": "google/gemma-2-9b-it"
+}
+
+# --- VARIÁVEIS GLOBAIS (CACHE NA VRAM) ---
+# Vamos guardar tudo na memória da H200
+loaded_models = {}
+loaded_tokenizers = {}
+
+def get_model_and_tokenizer(model_key):
+    """Carrega o modelo na VRAM apenas se ainda não estiver lá."""
+    global loaded_models, loaded_tokenizers
+    
+    if model_key not in loaded_models:
+        model_id = MODELS[model_key]
+        print(f"🐢 Cold Start: Carregando {model_id} na VRAM...")
+        
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16, # BF16 economiza memória e é rápido na H200
+            device_map="cuda"
+        )
+        
+        loaded_models[model_key] = model
+        loaded_tokenizers[model_key] = tokenizer
+        print(f"✅ {model_id} carregado e pronto!")
+        
+    return loaded_models[model_key], loaded_tokenizers[model_key]
+
+# --- FUNÇÃO DE GERAÇÃO (ZEROGPU) ---
+@spaces.GPU(duration=120) # 2 min é seguro para respostas longas de matemática
+def generate(message, history, model_selector):
+    # Identifica qual modelo o usuário quer
+    if "DeepSeek" in model_selector:
+        key = "deepseek"
+    elif "Llama" in model_selector:
+        key = "llama3"
+    elif "Gemma" in model_selector:
+        key = "gemma2"
+    else:
+        key = "deepseek" # Padrão
+    
+    model, tokenizer = get_model_and_tokenizer(key)
+    
+    # Formata o prompt (Cada modelo tem seu jeito, mas o tokenizer resolve)
+    # Convertendo histórico para formato de lista de dicts
+    messages = []
+    for user_msg, bot_msg in history:
+        if user_msg: messages.append({"role": "user", "content": user_msg})
+        if bot_msg: messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
+
+    # Aplica o template de chat correto para o modelo
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+    
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    # Gera a resposta
+    # Max tokens alto pq matemática exige passo-a-passo
+    outputs = model.generate(
+        **inputs, 
+        max_new_tokens=2048,
+        temperature=0.6,
+        do_sample=True
+    )
+    
+    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    return response
+
+# --- INTERFACE GRADIO ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🧮 Módulo Matemático & Lógico (H200)")
+    
+    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)"
+            ],
+            value="🐳 DeepSeek Math 7B (Especialista em Exatas)", 
+            label="Escolha o Especialista",
+            interactive=True
+        )
+    
+    chat = gr.ChatInterface(
+        fn=generate,
+        additional_inputs=[model_dropdown]
+    )
+
+if __name__ == "__main__":
+    demo.launch()