import spaces import gradio as gr import torch import transformers from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig import os title = """# Welcome to 🌟Tonic's✨StarCoder ✨StarCoder StarCoder2-15B model is a 15B parameter model trained on 600+ programming languages from The Stack v2, with opt-out requests excluded. The model uses Grouped Query Attention, a context window of 16,384 tokens with a sliding window attention of 4,096 tokens, and was trained using the Fill-in-the-Middle objective on 4+ trillion tokens. The model was trained with NVIDIA NeMo™ Framework using the NVIDIA Eos Supercomputer built with NVIDIA DGX H100 systems. You can build with this endpoint using✨StarCoder available here : [bigcode/starcoder2-15b](https://huggingface.co/bigcode/starcoder2-15b). You can also use ✨StarCoder 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:[MultiTransformer](https://huggingface.co/MultiTransformer) Math 🔍 [introspector](https://huggingface.co/introspector) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [SciTonic](https://github.com/Tonic-AI/scitonic)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗 """ model_path = "bigcode/starcoder2-3b" hf_token = os.getenv("HF_TOKEN") if not hf_token: raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.") tokenizer = AutoTokenizer.from_pretrained(model_path) quantization_config = BitsAndBytesConfig(load_in_8bit=True) model = AutoModelForCausalLM.from_pretrained( model_path, quantization_config=quantization_config) @spaces.GPU def generate_text(prompt, temperature=0.9, max_length=1200): # Encode the inputs inputs = tokenizer.encode(prompt, return_tensors="pt") attention_mask = torch.ones(inputs.shape, dtype=torch.long) inputs = inputs.to("cuda") attention_mask = attention_mask.to("cuda") outputs = model.generate( inputs, attention_mask=attention_mask, max_length=max_length, top_p=0.9, temperature=temperature, do_sample=True, pad_token_id=tokenizer.eos_token_id ) return tokenizer.decode(outputs[0]) def gradio_app(): with gr.Blocks() as demo: gr.Markdown(title) prompt = gr.Code(label="Enter your code prompt", value="def print_hello_world():") with gr.Row(): temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature") max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=200, label="Generate Length") generate_btn = gr.Button("Try✨StarCoder") output = gr.Code(label="✨StarCoder:", lines=40) generate_btn.click( fn=generate_text, inputs=[prompt, temperature, max_length], outputs=output ) demo.launch() if __name__ == "__main__": gradio_app()