import spaces import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer title = "# 👋🏻Welcome to🌟Tonic's⚖️StableCode2" description = """⚖️StableCode2 is a small sized coding llm that performs well in python ! You can also use [⚖️stabilityai/stable-code-3b](https://huggingface.co/stabilityai/stable-code-3b) by cloning this space. 🧬🔬🔍 Simply click here: Duplicate Space Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to 🌟 [EasyAGI](https://github.com/tonic-ai/EasyAGI) 🤗Big thanks to Ythe folks at huggingface for the ZeroGPU 🤗 To contribute to this space make a PR with a new example or cool new use-case for this one 🤗 """ tokenizer = AutoTokenizer.from_pretrained( "stabilityai/stable-code-3b", trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( "stabilityai/stable-code-3b", trust_remote_code=True, torch_dtype="auto", attn_implementation="flash_attention_2", ).to("cuda" if torch.cuda.is_available() else "cpu") @spaces.GPU def generate_code(prompt): inputs = tokenizer(prompt, return_tensors="pt").to(model.device) tokens = model.generate( **inputs, max_new_tokens=650, temperature=0.3, do_sample=False, ) generated_code = tokenizer.decode(tokens[0], skip_special_tokens=True) return generated_code with gr.Blocks() as demo: gr.Markdown(title) gr.Markdown(description) with gr.Row(): prompt = gr.Textbox(lines=2, placeholder="Enter your Python code prompt") output = gr.Textbox(label = "⚖️StableCode2") generate_button = gr.Button("Generate") generate_button.click(fn=generate_code, inputs=prompt, outputs=output) demo.launch()