Update app.py

app.py

@@ -5,12 +5,23 @@ import gradio as gr
 import requests
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from dotenv import load_dotenv
+import torch
 
-# ----------------- Load secrets -----------------
+# ----------------- Configuration & Secrets -----------------
+# Load environment variables from a .env file for local development
 load_dotenv()
-HF_TOKEN = os.environ.get("HF_TOKEN")              # Hugging Face Token
-SUPABASE_URL = os.environ.get("SUPABASE_URL")      # Supabase URL
-SUPABASE_KEY = os.environ.get("SUPABASE_KEY")      # Supabase API Key
+
+# --- Hugging Face ---
+HF_TOKEN = os.environ.get("HF_TOKEN")
+MODEL_NAME = "google/gemma-3-270m-it"
+
+# --- Supabase ---
+SUPABASE_URL = os.environ.get("SUPABASE_URL")
+SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
+
+# Check if essential environment variables are set
+if not all([HF_TOKEN, SUPABASE_URL, SUPABASE_KEY]):
+    raise ValueError("Missing required environment variables (HF_TOKEN, SUPABASE_URL, SUPABASE_KEY). Please check your .env file or environment settings.")
 
 HEADERS = {
     "apikey": SUPABASE_KEY,
@@ -18,52 +29,163 @@ HEADERS = {
     "Content-Type": "application/json"
 }
 
-# ----------------- Supabase functions -----------------
+# ----------------- Model Loading -----------------
+# Load the Gemma model and tokenizer from Hugging Face
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
+    # Load model with torch_dtype=torch.float16 for better performance if a GPU is available
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        token=HF_TOKEN,
+        torch_dtype=torch.float16,
+        device_map="auto" # Automatically use GPU if available
+    )
+    gen_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    print("✅ Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    gen_pipe = None
+
+# ----------------- Backend Logic (Supabase) -----------------
 def create_user(name: str):
+    """Creates a new user in Supabase and returns a unique API key."""
+    if not name or not name.strip():
+        return "⚠️ Error: User name cannot be empty.", ""
+
     api_key = str(uuid.uuid4())
-    data = {"name": name, "api_key": api_key}
-    response = requests.post(f"{SUPABASE_URL}/rest/v1/users",
-                             headers=HEADERS,
-                             data=json.dumps(data))
-    if response.status_code in [200, 201]:
-        return api_key
-    else:
-        return f"Erro: {response.text}"
-
-def check_key(api_key: str):
-    response = requests.get(f"{SUPABASE_URL}/rest/v1/users?api_key=eq.{api_key}",
-                            headers=HEADERS)
-    data = response.json()
-    return len(data) > 0
-
-# ----------------- Carregar modelo Gemma 270M -----------------
-model_name = "google/gemma-3-270m-it"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HF_TOKEN)
-gen_pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# ----------------- Função de geração de texto -----------------
-def generate_text(prompt, api_key):
+    data = {"name": name.strip(), "api_key": api_key}
+    
+    try:
+        response = requests.post(
+            f"{SUPABASE_URL}/rest/v1/users",
+            headers=HEADERS,
+            data=json.dumps(data)
+        )
+        response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx)
+        
+        if response.status_code in [200, 201]:
+            return f"✅ Success! Your API key has been created for user '{name}'.", api_key
+        else:
+            return f"⚠️ Error: Could not create user. Server responded with status {response.status_code}.", ""
+            
+    except requests.exceptions.RequestException as e:
+        return f"❌ Network Error: {e}", ""
+
+def check_key(api_key: str) -> bool:
+    """Checks if a given API key is valid by querying Supabase."""
+    if not api_key or not api_key.strip():
+        return False
+        
+    try:
+        response = requests.get(
+            f"{SUPABASE_URL}/rest/v1/users?api_key=eq.{api_key.strip()}",
+            headers=HEADERS
+        )
+        response.raise_for_status()
+        return len(response.json()) > 0
+    except requests.exceptions.RequestException:
+        return False
+
+# ----------------- Text Generation Logic -----------------
+def generate_text(prompt: str, api_key: str):
+    """Generates text using the Gemma model if the API key is valid."""
+    if gen_pipe is None:
+        return "❌ Model is not available. Please check the server logs."
+        
     if not check_key(api_key):
+        gr.Warning("API Key is invalid or missing! Please create a valid key.")
         return "API Key inválida!"
-    result = gen_pipe(prompt, max_length=100, do_sample=True)
-    return result[0]["generated_text"]
 
-# ----------------- Interface Gradio -----------------
-with gr.Blocks() as demo:
-    gr.Markdown("## Gemma 270M API - Hugging Face Space")
-    
-    with gr.Tab("Gerar Texto"):
-        prompt_input = gr.Textbox(label="Prompt")
-        api_key_input = gr.Textbox(label="API Key")
-        output_text = gr.Textbox(label="Resultado")
-        generate_btn = gr.Button("Gerar")
-        generate_btn.click(fn=generate_text, inputs=[prompt_input, api_key_input], outputs=output_text)
-
-    with gr.Tab("Criar API Key"):
-        name_input = gr.Textbox(label="Nome do Usuário")
-        key_output = gr.Textbox(label="Sua API Key")
-        create_btn = gr.Button("Criar")
-        create_btn.click(fn=create_user, inputs=name_input, outputs=key_output)
-
-demo.launch(server_name="0.0.0.0", server_port=7860)
+    if not prompt or not prompt.strip():
+        gr.Warning("Prompt cannot be empty.")
+        return ""
+
+    try:
+        # Generate text using the pipeline
+        result = gen_pipe(
+            prompt,
+            max_new_tokens=250, # Use max_new_tokens for more control
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95
+        )
+        return result[0]["generated_text"]
+    except Exception as e:
+        gr.Error(f"An error occurred during text generation: {e}")
+        return "Sorry, something went wrong."
+
+# ----------------- Gradio Interface -----------------
+# Custom CSS for a more polished look
+custom_css = """
+body { background-color: #F8F9FA !important; }
+.gradio-container { max-width: 800px !important; margin: auto !important; }
+footer { display: none !important; }
+.gr-button { background-color: #1d4ed8 !important; color: white !important; border-radius: 8px !important; }
+.gr-button:hover { background-color: #2563eb !important; }
+.gr-input { border-radius: 8px !important; }
+.gr-textbox { border-radius: 8px !important; }
+#title { text-align: center; color: #1e3a8a; font-size: 2.5rem; font-weight: 600; margin-bottom: 1rem; }
+#subtitle { text-align: center; color: #4b5563; font-size: 1.1rem; margin-bottom: 2rem; }
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
+    # Header
+    gr.Markdown("<h1 id='title'>✨ Gemma 270M API Service</h1>", elem_id="title")
+    gr.Markdown("<p id='subtitle'>Generate text with Google's Gemma model, secured with your own API key.</p>", elem_id="subtitle")
+
+    with gr.Tabs() as tabs:
+        # --- Text Generation Tab ---
+        with gr.TabItem("📝 Generate Text", id=0):
+            with gr.Column(gap="md"):
+                gr.Markdown("### Enter your prompt and API key below to generate text.")
+                api_key_input = gr.Textbox(
+                    label="🔑 API Key",
+                    placeholder="Enter your unique API key here...",
+                    type="password",
+                    lines=1
+                )
+                prompt_input = gr.Textbox(
+                    label="💬 Prompt",
+                    placeholder="e.g., Write a short story about a robot who discovers music...",
+                    lines=5
+                )
+                output_text = gr.Textbox(
+                    label="💡 Result",
+                    lines=8,
+                    interactive=False,
+                    placeholder="Generated text will appear here..."
+                )
+                with gr.Row():
+                    clear_btn_gen = gr.ClearButton(value="Clear", components=[prompt_input, output_text])
+                    generate_btn = gr.Button("Generate", variant="primary")
+
+        # --- API Key Creation Tab ---
+        with gr.TabItem("🔐 Create API Key", id=1):
+            with gr.Column(gap="md"):
+                gr.Markdown("### Need an API Key? Create one here.")
+                name_input = gr.Textbox(label="👤 Your Name", placeholder="Enter your name to associate with the key...")
+                key_status_output = gr.Textbox(label="Status", interactive=False)
+                key_output = gr.Textbox(label="✅ Your New API Key", interactive=False)
+                with gr.Row():
+                    clear_btn_key = gr.ClearButton(value="Clear", components=[name_input, key_status_output, key_output])
+                    create_btn = gr.Button("Create Key", variant="primary")
+
+    # --- Event Handlers ---
+    generate_btn.click(
+        fn=generate_text,
+        inputs=[prompt_input, api_key_input],
+        outputs=output_text,
+        api_name="generate_text"
+    )
+
+    create_btn.click(
+        fn=create_user,
+        inputs=name_input,
+        outputs=[key_status_output, key_output],
+        api_name="create_user"
+    )
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)