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 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)