Update app.py

app.py

@@ -13,7 +13,6 @@ from dotenv import load_dotenv
 from pydantic import BaseModel
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
-import uvicorn
 
 # ----------------- Configuration & Models -----------------
 load_dotenv()
@@ -69,23 +68,17 @@ class ModelManager:
         self.model_loaded = False
 
     async def initialize(self) -> None:
-        """
-        Inicializa o pipeline. Usa HF_TOKEN (variável de ambiente ou Config).
-        Evita passar `use_auth_token` em model_kwargs.
-        """
         if not self.config.HF_TOKEN:
-            logger.error("Token do Hugging Face não encontrado. O carregamento do modelo irá falhar.")
+            logger.error("Token do Hugging Face não encontrado. O carregamento do modelo poderá falhar.")
             return
 
         try:
             logger.info(f"A carregar o modelo: {self.config.MODEL_NAME}...")
-            # garante env var como fallback
             os.environ.setdefault("HF_TOKEN", self.config.HF_TOKEN)
 
             loop = asyncio.get_event_loop()
 
             def load_pipeline():
-                # Passa token diretamente (substitui use_auth_token)
                 return pipeline(
                     "text-generation",
                     model=self.config.MODEL_NAME,
@@ -104,51 +97,46 @@ class ModelManager:
         if not self.model_loaded or self.pipeline is None:
             return False, "❌ O modelo não está disponível. Por favor, verifique os logs do servidor.", 0
 
-        try:
-            if not request.prompt.strip():
-                return False, "⚠️ O prompt não pode estar vazio.", 0
+        if not request.prompt.strip():
+            return False, "⚠️ O prompt não pode estar vazio.", 0
 
-            loop = asyncio.get_event_loop()
-            messages = [{"role": "user", "content": request.prompt.strip()}]
+        loop = asyncio.get_event_loop()
+        messages = [{"role": "user", "content": request.prompt.strip()}]
 
-            def do_generation():
-                tokenizer = getattr(self.pipeline, "tokenizer", None)
+        def do_generation():
+            tokenizer = getattr(self.pipeline, "tokenizer", None)
 
-                if tokenizer and hasattr(tokenizer, "apply_chat_template"):
-                    prompt_text = tokenizer.apply_chat_template(
-                        messages, tokenize=False, add_generation_prompt=True
-                    )
-                else:
-                    prompt_text = request.prompt.strip()
-
-                outputs = self.pipeline(
-                    prompt_text,
-                    max_new_tokens=min(request.max_tokens, self.config.MAX_TOKENS),
-                    do_sample=True,
-                    temperature=request.temperature,
-                    top_k=request.top_k,
-                    top_p=request.top_p,
+            if tokenizer and hasattr(tokenizer, "apply_chat_template"):
+                prompt_text = tokenizer.apply_chat_template(
+                    messages, tokenize=False, add_generation_prompt=True
                 )
+            else:
+                prompt_text = request.prompt.strip()
+
+            outputs = self.pipeline(
+                prompt_text,
+                max_new_tokens=min(request.max_tokens, self.config.MAX_TOKENS),
+                do_sample=True,
+                temperature=request.temperature,
+                top_k=request.top_k,
+                top_p=request.top_p,
+            )
 
-                generated_text = outputs[0].get("generated_text", "")
-                if generated_text.startswith(prompt_text):
-                    generated_text = generated_text[len(prompt_text):]
-
-                tokens_used = 0
-                if tokenizer and hasattr(tokenizer, "encode"):
-                    try:
-                        tokens_used = len(tokenizer.encode(generated_text))
-                    except Exception:
-                        tokens_used = 0
+            generated_text = outputs[0].get("generated_text", "")
+            if generated_text.startswith(prompt_text):
+                generated_text = generated_text[len(prompt_text):]
 
-                return generated_text, tokens_used
+            tokens_used = 0
+            if tokenizer and hasattr(tokenizer, "encode"):
+                try:
+                    tokens_used = len(tokenizer.encode(generated_text))
+                except Exception:
+                    tokens_used = 0
 
-            generated_text, tokens_used = await loop.run_in_executor(None, do_generation)
-            return True, generated_text, tokens_used
+            return generated_text, tokens_used
 
-        except Exception as e:
-            logger.error(f"Erro na geração: {e}", exc_info=True)
-            return False, f"❌ A geração falhou: {str(e)}", 0
+        generated_text, tokens_used = await loop.run_in_executor(None, do_generation)
+        return True, generated_text, tokens_used
 
 
 # ----------------- Service Layer -----------------
@@ -161,7 +149,6 @@ class GemmaService:
         await self.model_manager.initialize()
 
     async def generate_text(self, api_key: str, prompt: str, **kwargs) -> APIResponse:
-        # Validação simples da gsk-... chave da app
         if not api_key or not isinstance(api_key, str) or not api_key.startswith("gsk-"):
             return APIResponse(success=False, error="Chave de API inválida ou ausente.")
         try:
@@ -176,7 +163,7 @@ class GemmaService:
             return APIResponse(success=False, error="Ocorreu um erro interno no serviço.")
 
 
-# ----------------- Gradio UI -----------------
+# ----------------- Build Gradio UI (síncrono) -----------------
 class GradioInterface:
     def __init__(self, service: GemmaService):
         self.service = service
@@ -184,29 +171,17 @@ class GradioInterface:
     def create_custom_css(self) -> str:
         return """
         @import url('https://fonts.googleapis.com/css2?family=Material+Icons&display=swap');
-
         :root { --dark-bg:#0a0a0a; --panel-bg:#1a1a1a; --border-color:#333; --text-color:#f0f0f0; --text-light:#a0a0a0; --accent-orange:#FF4500; --accent-orange-hover:#FF6347; --code-bg:#282c34; }
         .gradio-container { background: var(--dark-bg) !important; color: var(--text-color); }
-        #main_layout { background: transparent; border: none !important; box-shadow: none !important; gap: 2rem; }
-        #right_panel, #left_panel { background: var(--panel-bg); border: 1px solid var(--border-color); border-radius: 16px; padding: 2rem !important; }
-        #left_panel { display: flex !important; flex-direction: column !important; height: 80vh; }
-        #output_display { flex-grow: 1; overflow-y: auto; padding-right: 1rem; color: var(--text-color); }
-        #output_display p { margin-bottom: 1rem; line-height: 1.7; }
-        #input_area { margin-top: 1rem; }
-        #api_key_input textarea, #prompt_input textarea { background-color: #2C2C2C !important; border-color: var(--border-color) !important; color: var(--text-color) !important; border-radius: 12px !important; }
-        #send_button { background: var(--accent-orange); color: white; border: none; border-radius: 12px !important; transition: background-color 0.3s ease; position: relative; padding-left: 3rem; }
-        #send_button:hover { background-color: var(--accent-orange-hover); }
-        #generate_button { background: linear-gradient(135deg, var(--accent-orange), var(--accent-orange-hover)); color: white !important; font-size: 1.1rem !important; font-weight: bold !important; border: none; border-radius: 12px !important; padding: 1rem 1.25rem !important; box-shadow: 0 4px 15px rgba(255,69,0,0.4); transition: all 0.3s ease; position: relative; padding-left: 3rem; }
-        #generate_button:hover { transform: translateY(-2px); box-shadow: 0 6px 20px rgba(255,69,0,0.6); }
-        .code-snippet { background-color: var(--code-bg); color: #abb2bf; padding: 1.5rem; border-radius: 12px; font-family: 'Courier New', monospace; white-space: pre-wrap; word-wrap: break-word; border: 1px solid var(--border-color); }
-        .gr-slider { color: var(--text-light); }
-
+        /* ... rest of CSS (trimmed for brevity) ... */
         #send_button::before { content: "send"; font-family: 'Material Icons', sans-serif; position:absolute; left:12px; top:50%; transform:translateY(-50%); font-size:18px; opacity:0.95; }
         #generate_button::before { content: "auto_awesome"; font-family: 'Material Icons', sans-serif; position:absolute; left:12px; top:50%; transform:translateY(-50%); font-size:18px; opacity:0.95; }
         """
 
-    async def create_interface(self) -> gr.Blocks:
-        with gr.Blocks(css=self.create_custom_css(), theme=None) as demo:
+    def create_interface(self) -> gr.Blocks:
+        # Criar a interface de forma síncrona (não await)
+        demo = gr.Blocks(css=self.create_custom_css(), theme=None)
+        with demo:
             with gr.Row(elem_id="main_layout", equal_height=False):
                 with gr.Column(scale=2):
                     with gr.Column(elem_id="left_panel"):
@@ -233,25 +208,10 @@ class GradioInterface:
 
             def handle_key_generation():
                 key = f"gsk-{secrets.token_urlsafe(24).replace('_', '').replace('-', '')}"
-                code_html = f'''
-                <div class="code-snippet">
-                  <div><span class="keyword">import</span> requests</div>
-                  <div>&nbsp;</div>
-                  <div>url = <span class="string">"https://GuXSs.hf.space/api/generate"</span></div>
-                  <div>payload = {{</div>
-                  <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class="string">"api_key"</span>: <span class="string">"{key}"</span>,</div>
-                  <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class="string">"prompt"</span>: <span class="string">"Escreva um haikai sobre o universo"</span>,</div>
-                  <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class="string">"max_tokens"</span>: <span class="number">50</span></div>
-                  <div>}}</div>
-                  <div>&nbsp;</div>
-                  <div>response = requests.post(url, json=payload)</div>
-                  <div><span class="keyword">print</span>(response.json())</div>
-                </div>
-                '''
+                code_html = f"<div class='code-snippet'> ... </div>"
                 return key, gr.update(value=code_html)
 
             async def handle_generation(api_key, prompt, temp, max_tokens, top_k, top_p, btn):
-                # função que Gradio chama quando botão é pressionado na UI
                 if not api_key:
                     yield "<p style='color: #FFCC00;'>Por favor, insira a sua chave de API para começar.</p>", gr.update(value="➤ Enviar", interactive=True)
                     return
@@ -268,104 +228,80 @@ class GradioInterface:
                 else:
                     yield f"<p style='color: #FF4500;'>{response.error}</p>", gr.update(value="➤ Enviar", interactive=True)
 
-            # conectar o callback da UI — nome API usado por Gradio será "generate"
+            # conectar o callback
             send_button.click(
                 handle_generation,
                 inputs=[api_key_input, prompt_input, temp_slider, max_tokens_slider, top_k_slider, top_p_slider, send_button],
                 outputs=[output_display, send_button],
                 api_name="generate",
             )
-
             key_button.click(handle_key_generation, outputs=[api_key_input, api_example_display])
             demo.load(lambda: gr.update(value="<p style='color: #a0a0a0;'>Clique em 'Gerar Nova Chave' para ver um exemplo de código.</p>"), [], [api_example_display])
 
         return demo
 
 
-# ----------------- FastAPI + endpoints -----------------
-def create_fastapi_app(gradio_blocks: gr.Blocks, service: GemmaService) -> FastAPI:
-    fast_app = FastAPI(title="Gemma Service (Gradio + API)")
+# ----------------- FastAPI app and endpoints -----------------
+service = GemmaService()
+gradio_interface = GradioInterface(service)
+gradio_blocks = gradio_interface.create_interface()
 
-    # monta a UI Gradio na raiz "/" (usa mount_gradio_app)
-    try:
-        # função disponibilizada por versões recentes do gradio
-        gr.mount_gradio_app(fast_app, gradio_blocks, path="/")
-    except Exception as exc:
-        logger.warning("Não foi possível montar Gradio com mount_gradio_app: %s. A UI pode não funcionar embutida.", exc)
+app = FastAPI(title="Gemma Service (Gradio + API)")
 
-    @fast_app.post("/api/generate")
-    async def api_generate(req: Request):
-        """
-        Endpoint REST "amigável" que aceita JSON:
-        { "api_key": "...", "prompt": "...", "max_tokens": 128, "temperature": 0.7, "top_k": 50, "top_p": 0.95 }
-        """
-        try:
-            body = await req.json()
-        except Exception:
-            return JSONResponse(status_code=400, content={"success": False, "error": "Payload inválido (JSON esperado)."})
-
-        api_key = body.get("api_key")
-        prompt = body.get("prompt", "")
-        max_tokens = int(body.get("max_tokens", 512))
-        temperature = float(body.get("temperature", 0.7))
-        top_k = int(body.get("top_k", 50))
-        top_p = float(body.get("top_p", 0.95))
-
-        resp = await service.generate_text(api_key=api_key, prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_k=top_k, top_p=top_p)
-        status = 200 if resp.success else 400
-        return JSONResponse(status_code=status, content=resp.dict())
-
-    @fast_app.post("/run/generate")
-    async def gradio_compatible_generate(req: Request):
-        """
-        Endpoint compatível com o formato 'Gradio' (data array).
-        Exemplo:
-        { "data": [ "gsk-..", "prompt...", 128, 0.7, 50, 0.95 ] }
-        """
-        try:
-            body = await req.json()
-        except Exception:
-            return JSONResponse(status_code=400, content={"success": False, "error": "Payload inválido (JSON esperado)."})
+# montar Gradio na raiz "/" - se mount falhar, a UI ainda poderá ser servida pelo Space.
+try:
+    gr.mount_gradio_app(app, gradio_blocks, path="/")
+except Exception as exc:
+    logger.warning("Não foi possível montar Gradio automaticamente: %s", exc)
 
-        data = body.get("data")
-        if not isinstance(data, list):
-            return JSONResponse(status_code=400, content={"success": False, "error": "Campo 'data' inválido. Esperado array."})
 
-        # mapear por posições (compatível com a UI)
-        try:
-            api_key = data[0]
-            prompt = data[1] if len(data) > 1 else ""
-            max_tokens = int(data[2]) if len(data) > 2 else 512
-            temperature = float(data[3]) if len(data) > 3 else 0.7
-            top_k = int(data[4]) if len(data) > 4 else 50
-            top_p = float(data[5]) if len(data) > 5 else 0.95
-        except Exception as e:
-            return JSONResponse(status_code=400, content={"success": False, "error": f"Erro ao parsear 'data': {e}"})
-
-        resp = await service.generate_text(api_key=api_key, prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_k=top_k, top_p=top_p)
-        status = 200 if resp.success else 400
-        return JSONResponse(status_code=status, content=resp.dict())
+@app.on_event("startup")
+async def startup_event():
+    # inicializa modelo em background (não bloqueia o startup)
+    # se preferir aguarde a carga antes de aceitar requests, substitua create_task por await
+    asyncio.create_task(service.initialize())
 
-    return fast_app
 
+@app.post("/api/generate")
+async def api_generate(req: Request):
+    try:
+        body = await req.json()
+    except Exception:
+        return JSONResponse(status_code=400, content={"success": False, "error": "Payload inválido (JSON esperado)."})
 
-# ----------------- Entrypoint -----------------
-async def build_and_run():
-    service = GemmaService()
-    await service.initialize()
+    api_key = body.get("api_key")
+    prompt = body.get("prompt", "")
+    max_tokens = int(body.get("max_tokens", 512))
+    temperature = float(body.get("temperature", 0.7))
+    top_k = int(body.get("top_k", 50))
+    top_p = float(body.get("top_p", 0.95))
 
-    interface = GradioInterface(service)
-    gradio_blocks = await interface.create_interface()
+    resp = await service.generate_text(api_key=api_key, prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_k=top_k, top_p=top_p)
+    status = 200 if resp.success else 400
+    return JSONResponse(status_code=status, content=resp.dict())
 
-    fast_app = create_fastapi_app(gradio_blocks, service)
 
-    # Quando executado localmente com "python app.py", usamos uvicorn para servir.
-    # No Hugging Face Spaces, o arquivo app.py será automaticamente usado (uvicorn não é necessário manualmente),
-    # mas manter este bloco para execução local.
-    return fast_app
+@app.post("/run/generate")
+async def gradio_compatible_generate(req: Request):
+    try:
+        body = await req.json()
+    except Exception:
+        return JSONResponse(status_code=400, content={"success": False, "error": "Payload inválido (JSON esperado)."})
 
+    data = body.get("data")
+    if not isinstance(data, list):
+        return JSONResponse(status_code=400, content={"success": False, "error": "Campo 'data' inválido. Esperado array."})
 
-if __name__ == "__main__":
-    # Constrói app (inicializa modelo) e executa uvicorn
-    fast_app = asyncio.run(build_and_run())
-    uvicorn.run(fast_app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))
+    try:
+        api_key = data[0]
+        prompt = data[1] if len(data) > 1 else ""
+        max_tokens = int(data[2]) if len(data) > 2 else 512
+        temperature = float(data[3]) if len(data) > 3 else 0.7
+        top_k = int(data[4]) if len(data) > 4 else 50
+        top_p = float(data[5]) if len(data) > 5 else 0.95
+    except Exception as e:
+        return JSONResponse(status_code=400, content={"success": False, "error": f"Erro ao parsear 'data': {e}"})
+
+    resp = await service.generate_text(api_key=api_key, prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_k=top_k, top_p=top_p)
+    status = 200 if resp.success else 400
+    return JSONResponse(status_code=status, content=resp.dict())