Update app.py

app.py

@@ -1,169 +1,159 @@
 import os
-import uuid
-import json
+import secrets
+import html
 import asyncio
 import logging
-import time
-from datetime import datetime, timedelta
-from typing import Dict, List, Optional, Tuple, Any
 from dataclasses import dataclass
+from typing import Any, Optional, Tuple
 
 import gradio as gr
-import aiohttp
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from dotenv import load_dotenv
-from pydantic import BaseModel, ValidationError
-import secrets
-import plotly.graph_objects as go
-from plotly.subplots import make_subplots
+from transformers import pipeline
+from pydantic import BaseModel
 
 # ----------------- Configuration & Models -----------------
-load_dotenv()
 
 @dataclass
 class Config:
     HF_TOKEN: str = os.getenv("HF_TOKEN", "")
     MODEL_NAME: str = os.getenv("MODEL_NAME", "google/gemma-3-270m-it")
-    MAX_TOKENS: int = int(os.getenv("MAX_TOKENS", "1500"))
+    MAX_TOKENS: int = int(os.getenv("MAX_TOKENS", "2048"))
     LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO")
 
+
 class GenerationRequest(BaseModel):
     prompt: str
-    max_tokens: int = 500
-    temperature: float = 0.75
+    max_tokens: int = 512
+    temperature: float = 0.7
     top_k: int = 50
     top_p: float = 0.95
-    repetition_penalty: float = 1.1
+
 
 class APIResponse(BaseModel):
     success: bool
     data: Any = None
     error: Optional[str] = None
-    timestamp: datetime = datetime.now()
+
 
 # ----------------- Enhanced Logger -----------------
+
 def setup_logger():
     logging.basicConfig(
         level=getattr(logging, Config().LOG_LEVEL),
         format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-        handlers=[
-            logging.FileHandler('gemma_saas.log'),
-            logging.StreamHandler()
-        ]
+        handlers=[logging.FileHandler('gemma_saas.log'), logging.StreamHandler()]
     )
     return logging.getLogger(__name__)
 
+
 logger = setup_logger()
 
+
 # ----------------- Model Manager -----------------
 class ModelManager:
     def __init__(self, config: Config):
         self.config = config
-        self.tokenizer = None
-        self.model = None
         self.pipeline = None
         self.model_loaded = False
 
     async def initialize(self):
-        """Initialize the model, tokenizer, and pipeline asynchronously."""
         if not self.config.HF_TOKEN:
-            logger.error("Hugging Face token not found. Model loading will fail.")
-            self.model_loaded = False
+            logger.error("Token do Hugging Face não encontrado. O carregamento do modelo irá falhar.")
             return
+
         try:
-            logger.info(f"Loading model: {self.config.MODEL_NAME}...")
-            loop = asyncio.get_event_loop()
-
-            def load_components():
-                tokenizer = AutoTokenizer.from_pretrained(self.config.MODEL_NAME, token=self.config.HF_TOKEN)
-                model = AutoModelForCausalLM.from_pretrained(
-                    self.config.MODEL_NAME,
-                    token=self.config.HF_TOKEN,
+            logger.info(f"A carregar o modelo: {self.config.MODEL_NAME}...")
+            loop = asyncio.get_running_loop()
+
+            def load_pipeline():
+                # Use `use_auth_token` (aplicável em muitas versões do transformers)
+                return pipeline(
+                    task="text-generation",
+                    model=self.config.MODEL_NAME,
                     device_map="auto",
-                    torch_dtype="auto"
-                )
-                text_pipeline = pipeline(
-                    "text-generation",
-                    model=model,
-                    tokenizer=tokenizer,
+                    model_kwargs={"torch_dtype": "auto"},
+                    use_auth_token=self.config.HF_TOKEN,
                 )
-                return tokenizer, model, text_pipeline
 
-            self.tokenizer, self.model, self.pipeline = await loop.run_in_executor(None, load_components)
+            self.pipeline = await loop.run_in_executor(None, load_pipeline)
             self.model_loaded = True
-            logger.info("✅ Model loaded successfully!")
+            logger.info("✅ Modelo carregado com sucesso!")
         except Exception as e:
-            logger.error(f"❌ Error loading model: {e}")
+            logger.error(f"❌ Erro ao carregar o modelo: {e}")
             self.model_loaded = False
 
     async def generate(self, request: GenerationRequest) -> Tuple[bool, str, int]:
-        """Generate text based on the provided request."""
-        if not self.model_loaded:
+        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 len(request.prompt) > 8000:
-                return False, "⚠️ O prompt é muito longo (máximo de 8000 caracteres).", 0
 
-            loop = asyncio.get_event_loop()
-            
-            messages = [
-                {"role": "user", "content": request.prompt.strip()},
-            ]
+        if not request.prompt.strip():
+            return False, "⚠️ O prompt não pode estar vazio.", 0
+
+        try:
+            loop = asyncio.get_running_loop()
 
             def do_generation():
-                prompt = self.pipeline.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+                # Para a maioria dos modelos de geração textual, passamos o prompt diretamente
+                prompt_text = request.prompt.strip()
+
                 outputs = self.pipeline(
-                    prompt,
+                    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,
                 )
-                return outputs[0]["generated_text"][len(prompt):]
 
-            generated_text = await loop.run_in_executor(None, do_generation)
-            tokens_used = len(self.tokenizer.encode(generated_text))
+                # A saída típica é uma lista com dicionários contendo 'generated_text'
+                generated_text = outputs[0].get("generated_text", "")
+
+                # Contagem aproximada de tokens (usa o tokenizer do pipeline se disponível)
+                tokens_used = 0
+                try:
+                    tokenizer = getattr(self.pipeline, "tokenizer", None)
+                    if tokenizer is not None:
+                        # Evitar adicionar special tokens na contagem
+                        tokens_used = len(tokenizer.encode(generated_text, add_special_tokens=False))
+                    else:
+                        tokens_used = len(generated_text.split())
+                except Exception:
+                    tokens_used = len(generated_text.split())
+
+                return generated_text, tokens_used
+
+            generated_text, tokens_used = await loop.run_in_executor(None, do_generation)
             return True, generated_text, tokens_used
+
         except Exception as e:
-            logger.error(f"Generation error: {e}")
+            logger.error(f"Erro na geração: {e}")
             return False, f"❌ A geração falhou: {str(e)}", 0
 
+
 # ----------------- Service Layer -----------------
 class GemmaService:
     def __init__(self):
         self.config = Config()
         self.model_manager = ModelManager(self.config)
-        self._validate_config()
-
-    def _validate_config(self):
-        """Validate that required environment variables are set."""
-        if not self.config.HF_TOKEN:
-            raise ValueError("Missing required environment variable: HF_TOKEN")
 
     async def initialize(self):
         await self.model_manager.initialize()
 
-    async def generate_text(self, prompt: str, **kwargs) -> APIResponse:
-        """Generate text directly."""
+    async def generate_text(self, api_key: str, prompt: str, **kwargs) -> APIResponse:
+        if not api_key or not api_key.startswith("gsk-"):
+            return APIResponse(success=False, error="Chave de API inválida ou ausente.")
+
         try:
             request = GenerationRequest(prompt=prompt, **kwargs)
             success, text, tokens_used = await self.model_manager.generate(request)
-
             if success:
-                return APIResponse(
-                    success=True,
-                    data={"generated_text": text, "tokens_used": tokens_used}
-                )
+                return APIResponse(success=True, data={"generated_text": text, "tokens_used": tokens_used})
             else:
                 return APIResponse(success=False, error=text)
-
         except Exception as e:
-            logger.error(f"Service error during text generation: {e}")
+            logger.error(f"Erro de serviço durante a geração de texto: {e}")
             return APIResponse(success=False, error="Ocorreu um erro interno no serviço.")
 
+
 # ----------------- Enhanced UI -----------------
 class GradioInterface:
     def __init__(self, service: GemmaService):
@@ -172,139 +162,132 @@ class GradioInterface:
     def create_custom_css(self):
         return """
         :root {
-            --dark-bg: #111111;
-            --panel-bg: #1C1C1C;
-            --border-color: #333333;
-            --text-color: #E0E0E0;
-            --text-light: #A0A0A0;
-            --accent-orange: #FF4500;
-            --accent-orange-hover: #FF6347;
+            --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-color: var(--dark-bg) !important; }
-        #main_layout { background-color: transparent; border: none !important; box-shadow: none !important; }
-        #right_panel { background-color: var(--panel-bg); border-left: 1px solid var(--border-color); border-radius: 12px; padding: 2rem !important; }
-        #left_panel { background-color: var(--panel-bg); border-radius: 12px; padding: 1rem !important; display: flex !important; flex-direction: column !important; height: 70vh; }
-        #output_display { flex-grow: 1; overflow-y: auto; padding: 1rem; color: var(--text-color); }
-        #output_display p { margin-bottom: 1rem; line-height: 1.6; }
-        #prompt_row { border-top: 1px solid var(--border-color); padding-top: 1rem; }
-        #prompt_input textarea { background-color: #2C2C2C !important; border-color: var(--border-color) !important; color: var(--text-color) !important; border-radius: 8px !important; }
-        #send_button { background-color: var(--accent-orange); color: white; border: none; border-radius: 50% !important; width: 50px !important; height: 50px !important; min-width: 50px !important; transition: background-color 0.3s ease; }
+        .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; }
         #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.2rem !important;
-            font-weight: bold !important;
-            border: none;
-            border-radius: 12px !important;
-            padding: 1rem !important;
-            box-shadow: 0 4px 15px rgba(255, 69, 0, 0.4);
-            transition: all 0.3s ease;
-        }
-        #generate_button:hover {
-            transform: translateY(-2px);
-            box-shadow: 0 6px 20px rgba(255, 69, 0, 0.6);
+            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 !important; box-shadow: 0 4px 15px rgba(255, 69, 0, 0.4); transition: all 0.3s ease;
         }
-        .gr-label { color: var(--text-light) !important; }
-        h2 { color: white; border-bottom: 1px solid var(--border-color); padding-bottom: 0.5rem; margin-bottom: 1rem; }
-        #info_text { color: var(--text-light); line-height: 1.7; }
+        #generate_button:hover { transform: translateY(-2px); box-shadow: 0 6px 20px rgba(255, 69, 0, 0.6); }
+        h2, h3 { color: white; border-bottom: 1px solid var(--border-color); padding-bottom: 0.75rem; margin-bottom: 1.5rem; font-weight: 600; }
+        .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); }
+        .code-snippet .keyword { color: #c678dd; } .code-snippet .string { color: #98c379; } .code-snippet .number { color: #d19a66; }
+        .gr-slider { color: var(--text-light); }
         """
 
-    async def create_interface(self):
+    def create_interface(self):
         with gr.Blocks(css=self.create_custom_css(), theme=None) as app:
             with gr.Row(elem_id="main_layout", equal_height=False):
-                with gr.Column(scale=2, elem_id="left_panel_container"):
+                with gr.Column(scale=2):
                     with gr.Column(elem_id="left_panel"):
-                        output_display = gr.Markdown(elem_id="output_display", value="<p style='color: #A0A0A0;'>Sua resposta aparecerá aqui...</p>")
-                        with gr.Row(elem_id="prompt_row"):
-                            prompt_input = gr.Textbox(
-                                show_label=False,
-                                placeholder="Digite sua mensagem aqui...",
-                                elem_id="prompt_input",
-                                scale=10
-                            )
-                            send_button = gr.Button("➤", elem_id="send_button", scale=1)
-
-                with gr.Column(scale=1, elem_id="right_panel"):
-                    gr.Markdown("## Informações")
-                    gr.Markdown(
-                        """
-                        Este é um ambiente interativo para o modelo de linguagem **Gemma**.
-                        
-                        - **Como usar:** Digite seu prompt na caixa de texto à esquerda e clique no botão de envio para gerar uma resposta.
-                        - **Gerar Chave:** Use o botão abaixo para gerar uma chave de API de exemplo.
-                        """,
-                        elem_id="info_text"
-                    )
-                    key_button = gr.Button("✨ Gerar Key", elem_id="generate_button")
-                    key_display = gr.Markdown()
-
-
-            # --- Event Handlers ---
-            async def handle_generation(prompt):
+                        output_display = gr.Markdown(elem_id="output_display", value="<p style='color: #a0a0a0;'>A sua resposta aparecerá aqui...</p>")
+                        with gr.Column(elem_id="input_area"):
+                            api_key_input = gr.Textbox(label="A Sua Chave de API", placeholder="Cole a sua chave gsk-... aqui", type="password", elem_id="api_key_input")
+                            with gr.Row():
+                                prompt_input = gr.Textbox(show_label=False, placeholder="Digite a sua mensagem...", elem_id="prompt_input", scale=10)
+                                send_button = gr.Button("➤ Enviar", elem_id="send_button", scale=2)
+
+                with gr.Column(scale=1):
+                    with gr.Column(elem_id="right_panel"):
+                        gr.Markdown("## Controlo")
+                        key_button = gr.Button("✨ Gerar Nova Chave", elem_id="generate_button")
+
+                        with gr.Accordion("Parâmetros Avançados", open=False):
+                            temp_slider = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperatura")
+                            max_tokens_slider = gr.Slider(minimum=64, maximum=self.service.config.MAX_TOKENS, value=512, step=64, label="Max Tokens")
+                            top_k_slider = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-K")
+                            top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+
+                        gr.Markdown("### Como Usar a API")
+                        api_example_display = gr.HTML("<p style='color: #a0a0a0;'>Clique em 'Gerar Nova Chave' para ver um exemplo de código.</p>")
+
+            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://SEU_SPACE.hf.space/run/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>
+                """
+
+                return gr.Textbox.update(value=key, interactive=True), api_example_display.update(value=code_html)
+
+            async def handle_generation(api_key, prompt, temp, max_tokens, top_k, top_p, btn):
+                if not api_key:
+                    yield "<p style='color: #FFCC00;'>Por favor, insira a sua chave de API para começar.</p>", gr.Button.update(value="➤ Enviar", interactive=True)
+                    return
                 if not prompt:
-                    # FIX: Use yield and return to exit the generator correctly
-                    yield "<p style='color: #FFCC00;'>Por favor, digite um prompt para começar.</p>"
+                    yield "<p style='color: #FFCC00;'>Por favor, digite um prompt.</p>", gr.Button.update(value="➤ Enviar", interactive=True)
                     return
-                
-                # Show a loading indicator
-                yield "<p style='color: #A0A0A0;'>Gerando resposta...</p>"
 
-                response = await self.service.generate_text(prompt=prompt)
-                
-                if response.success:
-                    yield response.data["generated_text"]
-                else:
-                    yield f"<p style='color: #FF4500;'>{response.error}</p>"
+                yield "<p style='color: #a0a0a0;'>A gerar resposta...</p>", gr.Button.update(value="A gerar...", interactive=False)
 
-            def handle_key_generation():
-                """Generates a random API key in the specified format."""
-                random_part = secrets.token_urlsafe(24).replace("_", "").replace("-", "")
-                key = f"gsk-{random_part}"
-                return f"<p style='color: #A0A0A0; text-align: center; margin-top: 1rem;'>Sua chave de exemplo:</p><pre style='background: #2C2C2C; padding: 1rem; border-radius: 8px; text-align: center; word-wrap: break-word;'><code>{key}</code></pre>"
+                # chama o serviço de geração
+                response = await self.service.generate_text(
+                    api_key=api_key,
+                    prompt=prompt,
+                    temperature=float(temp),
+                    max_tokens=int(max_tokens),
+                    top_k=int(top_k),
+                    top_p=float(top_p),
+                )
 
+                if response.success:
+                    formatted_text = html.escape(response.data["generated_text"]).replace("\n", "<br>")
+                    yield formatted_text, gr.Button.update(value="➤ Enviar", interactive=True)
+                else:
+                    yield f"<p style='color: #FF4500;'>{response.error}</p>", gr.Button.update(value="➤ Enviar", interactive=True)
 
-            # --- Wiring ---
-            key_button.click(
-                handle_key_generation,
-                inputs=[],
-                outputs=[key_display]
-            )
             send_button.click(
                 handle_generation,
-                inputs=[prompt_input],
-                outputs=[output_display]
-            )
-            prompt_input.submit(
-                handle_generation,
-                inputs=[prompt_input],
-                outputs=[output_display]
+                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])
+
         return app
 
+
 # ----------------- Main Application -----------------
-async def main():
-    """Main application entry point"""
+
+def main():
     try:
         service = GemmaService()
-        await service.initialize()
-        
+        # inicializa o modelo (bloqueante, mas necessário antes de lançar a UI)
+        asyncio.run(service.initialize())
+
         interface = GradioInterface(service)
-        app = await interface.create_interface()
-        
-        app.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            share=False, 
-            debug=False,
-            show_error=True
-        )
+        app = interface.create_interface()
+
+        # Lança a aplicação Gradio (bloqueia até terminar)
+        app.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=False, show_error=True)
     except Exception as e:
-        logger.critical(f"Failed to start application: {e}", exc_info=True)
-        raise
+        logger.critical(f"Falha ao iniciar a aplicação: {e}", exc_info=True)
+
 
 if __name__ == "__main__":
-    # To run this, you need a .env file with:
-    # HF_TOKEN="your_hugging_face_token"
-    asyncio.run(main())
+    main()