app.py

from transformers import pipeline, AutoTokenizer
import gradio as gr
import spaces
import os
import time
from collections import defaultdict

# Hidden state for the counter
# counter_state = gr.State(0)
hf_token = os.getenv("HF_TOKEN")

model_name = "markrodrigo/Qwen-3.5-2B-Spatial-SQL-1.1"   # or any model / local path

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    trust_remote_code=True,          # needed for many Qwen models
)

pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer, device_map="auto", token=hf_token)

usage_tracker = defaultdict(list)
MAX_REQUESTS_PER_DAY = 3
TRUSTED_USERS = {"markrodrigo"}   # ← Add your username here

# The Alpaca instruction prompt format
'''ALPACA_TEMPLATE = """<|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a helpful assistant. You are an expert at PostGIS and Postgresql and SQL and psql. <|eot_id|><|start_header_id|>user<|end_header_id|>
### Instruction: Write a PostGIS SQL statement for the following.
{instruction}
### Input:
{input}
### Response:
<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""'''

# Define your list of pre-set example prompts
PRESET_EXAMPLES = [
    "What is the area for the polygon? : 'Polygon ((-3.7515154 40.3855551, -3.7514972 40.3856581, -3.7507005 40.3855767, -3.7507167 40.3854722, -3.7515154 40.3855551))'",
    "What is the centroid for the polygon? : 'Polygon ((-3.6934636 40.4808785, -3.6933352 40.4811486, -3.6930125 40.4810598, -3.693141 40.4807897, -3.6934636 40.4808785))'",
    "What is the thousand meter buffer for the following point? : 'Point(-8.7522658 41.3862664)'",
    "How long is the line? : 'LINESTRING (-3.6976693 40.4263178, -3.6986082 40.4258729)'",
    "How far apart is the point and line? : 'Point(-109.87549823 38.60574249)' 'LineString(-109.24324628 38.76349931, -109.4821773 38.6875815)'"
]

# Custom CSS targeting the button by its ID
custom_css = """
#special-button {
    background-color: #0052FF !important;   /* Your desired color */
    color: white !important;
    border: none !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
}
#special-button:hover {
    background-color: #CAD3DE !important;
    transform: scale(1.03);
    transition: all 0.2s ease;
}
#clear-chat-btn {
    width: 200px !important;
    min-width: 200px !important;
    max-width: 200px !important;
    height: 42px !important;
    background-color: #FF5733 !important;   /* Main hex color */
    color: #FFFFFF !important;              /* White text */
    border: none !important;
    font-size: 14px !important;
    font-weight: 600 !important;
    padding: 8px 16px !important;
    border-radius: 8px !important;          /* Optional rounded look */
}
#clear-chat-btn:hover {
    background-color: #C70039 !important;   /* Hover hex color (darker) */
    transform: translateY(-1px);
    box-shadow: 0 4px 12px rgba(199, 0, 57, 0.3) !important;
}
"""
global_counter = 0

# def get_latest_counter():
    # global global_counter
    # return global_counter  # Dynamically fetches the current server value

def increment_counter():
    global global_counter
    global_counter += 1
    return global_counter 

@spaces.GPU
def respond(user_message, chat_history, profile: gr.OAuthProfile | None):  # , counter):
        # === BLOCK ANONYMOUS USERS ===
    if profile is None:
        raise gr.Error("You must be logged in with a Hugging Face account to use this demo.")
    
    username = profile.username
    print("Debug username:", username)   # You can keep this temporarily
    
    now = time.time()
    usage_tracker[username] = [ts for ts in usage_tracker[username] if now - ts < 86400]
    
    # Trusted users (you) get unlimited
    if username in TRUSTED_USERS:
        pass
    else:
        if len(usage_tracker[username]) >= MAX_REQUESTS_PER_DAY:
            raise gr.Error(f"Daily limit reached ({MAX_REQUESTS_PER_DAY} requests/day).")
    
    usage_tracker[username].append(now)

    # global global_counter
    # global_counter += 1


    # 1. Increment counter
    # counter += 1

    chat_history = chat_history or []

    if not user_message or not user_message.strip():
        return chat_history, ""

    # Modern format for Gradio 5/6
    chat_history.append({"role": "user", "content": user_message})
    chat_history.append({"role": "assistant", "content": None})

    # prompt = ALPACA_TEMPLATE.format(instruction=user_message, input="")
    # 1. Build the messages
    '''messages = [
        {"role": "system", "content": "Write a PostGIS SQL statement for the following."},
        {"role": "user", "content": "What is the thousand meter buffer for the following point? : 'Point(-3.78621945 40.4463195)'"},
    ]'''

    messages = [
        {"role": "system", "content": "Write a PostGIS SQL statement for the following."},
        {"role": "user", "content": user_message},
    ]

    # 2. Apply ChatML template
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,   # very important → adds <|im_start|>assistant\n
        enable_thinking=False,        # disable <think> tags (Qwen3)
    )

    sequences = pipe(
        prompt,
        max_new_tokens=256,
        return_full_text=False,
        temperature=0.4,
        top_k=100,
        do_sample=True,
    )

    bot_response = sequences[0]["generated_text"].strip()
    chat_history[-1]["content"] = bot_response

    return chat_history, "", global_counter  # , counter



with gr.Blocks(title="Text to PostGIS Postgresql via Qwen 3.5") as demo:
    gr.LoginButton()
    gr.Markdown("# Natural Language to Spatial SQL.\n###  Convert natural language and spatial primitives to PostGIS with Qwen 3.5")
    chatbot = gr.Chatbot(
        label="Chat",
        height=400,
        # type="messages",
        # show_copy_button=True,
    )

    with gr.Row():
        with gr.Column(scale=5):
            msg = gr.Textbox(
                placeholder="Natural Language : WKT format",
                lines=2,
                container=False
            )
        with gr.Column(scale=1, min_width=100):
            submit_btn = gr.Button("Submit", variant="primary", elem_id="special-button" )  # gr.themes.Ocean()

    gr.Markdown("### Quick Examples")
    gr.Examples(
        examples=PRESET_EXAMPLES,
        inputs=msg,
        label="Click an example → then click Submit"
    )
    # counter_display = gr.Textbox(label="Request Counter:", value=0, interactive=False)
    # submit_btn.click(fn=respond, inputs=[msg, chatbot, counter_state], outputs=[chatbot, msg, counter_state, counter_display])
    # msg.submit(fn=respond, inputs=[msg, chatbot, counter_state], outputs=[chatbot, msg, counter_state, counter_display])


    clear_btn = gr.Button("Clear Chat", elem_id="clear-chat-btn")
    clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg])
    gr.Markdown("## 🚀 Space Request Tracker")
    # gr.Textbox(value=str(global_counter), interactive=False)
    counter_display = gr.Number(value=global_counter, label="Global Counter")

    submit_btn.click(fn=respond, inputs=[msg, chatbot], outputs=[chatbot, msg, counter_display]).then(
        fn=increment_counter, 
        inputs=[], 
        outputs=counter_display
    )
    msg.submit(fn=respond, inputs=[msg, chatbot], outputs=[chatbot, msg, counter_display])
    # Display counter (optional)

    # CRITICAL: This triggers every time the page loads or refreshes!
    # demo.load(fn=get_latest_counter, inputs=[], outputs=counter_display)
    demo.load(fn=lambda: global_counter, inputs=[], outputs=counter_display)
    

if __name__ == "__main__":
    print("Gradio version:", gr.__version__)
    demo.launch(css=custom_css)