Create app.py

app.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import re
+
+import spaces
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+MODEL_ID = "prem-research/prem-1B-SQL"
+MODEL_REVISION = "44dd7fcf9227af4efed936bf29323c61bf66aad1"
+
+MAX_INPUT_TOKENS = 3072
+DEFAULT_MAX_NEW_TOKENS = 256
+MAX_NEW_TOKENS = 512
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    revision=MODEL_REVISION,
+    use_fast=True,
+)
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# ZeroGPU requires the GPU model to be placed on CUDA at module load time.
+model = (
+    AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        revision=MODEL_REVISION,
+        dtype=torch.float16,
+        use_safetensors=True,
+    )
+    .to("cuda")
+    .eval()
+)
+
+SQL_START = re.compile(r"\b(?:WITH|SELECT)\b", re.IGNORECASE)
+CODE_BLOCK = re.compile(
+    r"```(?:sql)?\s*(.*?)```",
+    re.IGNORECASE | re.DOTALL,
+)
+
+EXAMPLE_PROMPT = """Write exactly one read-only SQLite query.
+Use only the tables and columns in the schema.
+Return SQL only, without Markdown or explanation.
+
+Database schema:
+
+CREATE TABLE customers (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    city TEXT
+);
+
+Question: How many customers are there?
+
+SQL:
+"""
+
+
+def extract_sql(generated: str) -> str:
+    text = generated.strip()
+
+    fenced = CODE_BLOCK.search(text)
+    if fenced:
+        text = fenced.group(1).strip()
+
+    start = SQL_START.search(text)
+    if start:
+        text = text[start.start():]
+
+    return text.removesuffix("```").strip()
+
+
+@spaces.GPU(duration=45)
+def generate_sql(prompt: str, max_new_tokens: int) -> str:
+    prompt = (prompt or "").strip()
+
+    if not prompt:
+        raise gr.Error("Prompt cannot be empty.")
+
+    token_limit = max(
+        32,
+        min(MAX_NEW_TOKENS, int(max_new_tokens or DEFAULT_MAX_NEW_TOKENS)),
+    )
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        add_special_tokens=True,
+        truncation=False,
+    )
+
+    input_length = int(inputs["input_ids"].shape[-1])
+    if input_length > MAX_INPUT_TOKENS:
+        raise gr.Error(
+            f"Prompt has {input_length:,} tokens; "
+            f"the maximum is {MAX_INPUT_TOKENS:,}."
+        )
+
+    inputs = {
+        name: tensor.to("cuda")
+        for name, tensor in inputs.items()
+    }
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=token_limit,
+            do_sample=False,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+            use_cache=True,
+        )
+
+    new_tokens = output[0, input_length:]
+    generated = tokenizer.decode(
+        new_tokens,
+        skip_special_tokens=True,
+    )
+
+    sql = extract_sql(generated)
+    if not sql:
+        raise gr.Error("The model returned an empty response.")
+
+    return sql
+
+
+with gr.Blocks(title="Prem 1B SQL API") as demo:
+    gr.Markdown(
+        """
+        # Prem 1B SQL
+
+        Generate SQL from a schema-aware natural-language prompt.
+        The calling application must validate all generated SQL.
+        """
+    )
+
+    prompt_input = gr.Textbox(
+        value=EXAMPLE_PROMPT,
+        label="Schema-aware prompt",
+        lines=18,
+    )
+
+    max_tokens_input = gr.Slider(
+        minimum=32,
+        maximum=MAX_NEW_TOKENS,
+        value=DEFAULT_MAX_NEW_TOKENS,
+        step=1,
+        label="Maximum output tokens",
+    )
+
+    generate_button = gr.Button("Generate SQL", variant="primary")
+    sql_output = gr.Code(label="Generated SQL", language="sql")
+
+    generate_button.click(
+        fn=generate_sql,
+        inputs=[prompt_input, max_tokens_input],
+        outputs=sql_output,
+        api_name="generate",
+        concurrency_limit=1,
+    )
+
+demo.queue(max_size=20, default_concurrency_limit=1)
+
+if __name__ == "__main__":
+    demo.launch()