import spaces import gradio as gr import torch from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer, ) from threading import Thread from typing import Generator # --------------------------------------------------------------------------- # Module-scope model loading - ZeroGPU manages GPU offload transparently # --------------------------------------------------------------------------- MODEL_ID = "Qwen/Qwen3-Coder-30B-A3B-Instruct" quant_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16, ) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, quantization_config=quant_config, device_map="auto", torch_dtype=torch.bfloat16, trust_remote_code=True, ) model.eval() DEFAULT_SYSTEM = "You are an expert coding assistant. Write clean, efficient, well-documented code." # --------------------------------------------------------------------------- # ZeroGPU-decorated generation - xlarge for 30B MoE model # --------------------------------------------------------------------------- @spaces.GPU(duration=300) def generate( messages: list[dict], temperature: float, top_p: float, max_new_tokens: int, ) -> str: """Run model inference inside a ZeroGPU worker process. Args are pickled across the process boundary. Returns CPU text - safe for unpickling in the main process. """ inputs = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", ).to(model.device) with torch.inference_mode(): outputs = model.generate( inputs, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, do_sample=temperature > 0.0, pad_token_id=tokenizer.eos_token_id, ) generated = outputs[0][inputs.shape[1]:] return tokenizer.decode(generated, skip_special_tokens=True) # --------------------------------------------------------------------------- # Streaming variant - yields tokens as they're generated # --------------------------------------------------------------------------- @spaces.GPU(duration=300) def generate_stream( messages: list[dict], temperature: float, top_p: float, max_new_tokens: int, ) -> Generator[str, None, None]: """Stream tokens from the model one-by-one.""" inputs = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", ).to(model.device) streamer = TextIteratorStreamer( tokenizer, skip_prompt=True, skip_special_tokens=True, ) generation_kwargs = dict( inputs=inputs, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, do_sample=temperature > 0.0, pad_token_id=tokenizer.eos_token_id, streamer=streamer, ) thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() for token in streamer: yield token # --------------------------------------------------------------------------- # Non-streaming wrapper (for API endpoint) # --------------------------------------------------------------------------- def predict( message: str, history: list, system_prompt: str, temperature: float, top_p: float, max_tokens: int, ): """Chat function - called both from UI and the auto-generated Gradio API.""" messages = [{"role": "system", "content": system_prompt}] for user_msg, asst_msg in history: messages.append({"role": "user", "content": user_msg}) if asst_msg: messages.append({"role": "assistant", "content": asst_msg}) messages.append({"role": "user", "content": message}) output = generate(messages, temperature, top_p, max_tokens) return output # --------------------------------------------------------------------------- # Streaming chat handler # --------------------------------------------------------------------------- def chat_fn( message: str, history: list, system_prompt: str, temperature: float, top_p: float, max_tokens: int, ): """Generator that yields partial (message, history) tuples for streaming UI.""" messages = [{"role": "system", "content": system_prompt}] for user_msg, asst_msg in history: messages.append({"role": "user", "content": user_msg}) if asst_msg: messages.append({"role": "assistant", "content": asst_msg}) messages.append({"role": "user", "content": message}) partial = "" for token in generate_stream(messages, temperature, top_p, max_tokens): partial += token yield partial # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- LANGUAGES = ["python", "javascript", "typescript", "rust", "go", "java", "cpp", "csharp", "ruby", "php", "sql", "bash", "html", "css", "json", "yaml"] def build_examples(): return [ ["Write a Python async function that downloads a URL and retries 3 times on failure."], ["Create a Rust function that reads a CSV file and returns the row count."], ["Explain the difference between an interface and a type in TypeScript with examples."], ["Write a Go HTTP server that serves static files on port 8080 with CORS support."], ["Refactor this Python class to use dependency injection: class Database: ..."], ] # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- def create_ui(): with gr.Blocks( title="CodeCraft - AI Coding Assistant", theme=gr.themes.Soft( primary_hue="indigo", neutral_hue="slate", ), fill_width=True, ) as demo: gr.Markdown( "# CodeCraft - AI Coding Assistant\n" "Powered by **Qwen3-Coder-30B-A3B-Instruct** (MoE, 3B active) - ZeroGPU xlarge" ) chatbot = gr.Chatbot( label="Conversation", placeholder="Ask me anything about code...", render_markdown=True, show_copy_button=True, height=500, ) with gr.Row(): msg = gr.Textbox( label="Your message", placeholder="Write a Python async function that downloads a URL...", scale=8, container=False, ) submit_btn = gr.Button("Send", variant="primary", scale=1, min_width=80) clear_btn = gr.Button("Clear", scale=1, min_width=80) with gr.Accordion("Settings", open=False): with gr.Row(): system_prompt = gr.Textbox( label="System Prompt", value=DEFAULT_SYSTEM, lines=2, scale=3, ) with gr.Column(scale=1): temperature = gr.Slider( label="Temperature", minimum=0.0, maximum=1.5, value=0.3, step=0.05, ) with gr.Row(): top_p = gr.Slider( label="Top-P", minimum=0.6, maximum=1.0, value=0.9, step=0.05, ) max_tokens = gr.Slider( label="Max Tokens", minimum=128, maximum=8192, value=2048, step=128, ) gr.Examples( examples=build_examples(), inputs=[msg], label="Try these prompts", ) # -- State: chat history -- history_state = gr.State([]) # -- Event wiring -- def respond(message, history, system, temp, top_p_val, max_tok): if not message.strip(): return "", history, history history = history + [(message, None)] yield "", history, [] for partial in chat_fn(message, history[:-1], system, temp, top_p_val, max_tok): history[-1] = (message, partial) yield "", history, [] yield "", history, [message] msg.submit( respond, inputs=[msg, history_state, system_prompt, temperature, top_p, max_tokens], outputs=[msg, chatbot, history_state], concurrency_limit=4, api_name="predict", ) submit_btn.click( respond, inputs=[msg, history_state, system_prompt, temperature, top_p, max_tokens], outputs=[msg, chatbot, history_state], concurrency_limit=4, api_name=False, ) def clear_conversation(): return [], "", [] clear_btn.click( clear_conversation, outputs=[history_state, chatbot, msg], concurrency_limit=4, ) gr.Markdown( """ ### API This Space exposes a REST API at `/gradio_api/call/predict`. See the [Gradio docs](https://www.gradio.app/guides/sharing-your-app#api) for usage. """ ) return demo if __name__ == "__main__": demo = create_ui() demo.queue(default_concurrency_limit=4) demo.launch()