changing to starcoder

README.md

@@ -1,8 +1,8 @@
 ---
-title: Neo
-emoji: 🐇🥷🏻
-colorFrom: green
-colorTo: gray
+title: StarCoder2
+emoji: ✨2️⃣✨
+colorFrom: yellow
+colorTo: blue
 sdk: gradio
 sdk_version: 4.19.0
 app_file: app.py

app.py

@@ -5,15 +5,14 @@ import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
-title = """# Welcome to 🌟Tonic's🐇🥷🏻Neo
-WhiteRabbit🐇🥷🏻Neo is a model series that can be used for offensive and defensive cybersecurity. You can build with this endpoint using🐇🥷🏻Neo available here : [WhiteRabbitNeo/WhiteRabbitNeo-33B-v1.5](https://huggingface.co/WhiteRabbitNeo/WhiteRabbitNeo-33B-v1.5). You can also use 🐇🥷🏻Neo by cloning this space. Simply click here: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/neo?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
+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: <a style="display:inline-block" href="https://huggingface.co/spaces/Tonic/starcoder2?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></h3> 
 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 🤗
 """
 
-
 default_system_prompt = """SYSTEM: You are an AI that code. Answer with code."""
 
-model_path = "WhiteRabbitNeo/WhiteRabbitNeo-33B-v1.5"
+model_path = "bigcode/starcoder2-15b"
 
 
 hf_token = os.getenv("HF_TOKEN")
@@ -22,61 +21,43 @@ if not hf_token:
 
 model = AutoModelForCausalLM.from_pretrained(
     model_path,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
     device_map="auto",
-    load_in_4bit=True,
-    load_in_8bit=False,
-    trust_remote_code=True,
+#   trust_remote_code=True,
     token=hf_token,
 )
 
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 
-@spaces.GPU
-def generate_text(custom_prompt, user_input, temperature, generate_len, top_p, top_k):
-    system_prompt = custom_prompt if custom_prompt else default_system_prompt
-    llm_prompt = f"{system_prompt} \nUSER: {user_input} \nASSISTANT: "
-    
-    tokens = tokenizer.encode(llm_prompt, return_tensors="pt")
-    tokens = tokens.to("cuda")
+# import gradio as gr
+# from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
-    length = tokens.shape[1]
-    with torch.no_grad():
-        output = model.generate(
-            input_ids=tokens,
-            max_length=length + generate_len,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            num_return_sequences=1,
-        )
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    answer = generated_text[len(llm_prompt):].strip()
-    
-    return answer
+# checkpoint = "bigcode/starcoder2-15b"
+# quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+# tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+# model = AutoModelForCausalLM.from_pretrained(checkpoint, quantization_config=quantization_config).to("cuda")
+
+def generate_text(prompt, temperature, max_length):
+    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(inputs, max_length=max_length, top_p=0.9, temperature=temperature)
+    return tokenizer.decode(outputs[0])
 
 def gradio_app():
     with gr.Blocks() as demo:
         gr.Markdown(title)
+        output = gr.Code(label="Generated Code", lines=40)
         with gr.Row():
-            custom_prompt = gr.Textbox(label="🐇🥷🏻NeoCustom System Prompt (optional)", placeholder="Leave blank to use the default prompt...")
-            instruction = gr.Textbox(label="Your Instruction", placeholder="Type your question here...")
+            generate_btn = gr.Button("Generate")
         with gr.Row():
             temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
-            generate_len = gr.Slider(minimum=100, maximum=1024, step=10, value=100, label="Generate Length")
-            top_p = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=1.0, label="Top P")
-            top_k = gr.Slider(minimum=0, maximum=100, step=1, value=50, label="Top K")
+            max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=100, label="Generate Length")
         with gr.Row():
-            generate_btn = gr.Button("Generate")
-        output = gr.Code(label="🐇🥷🏻Neo:", lines=10)
+            prompt = gr.Textbox(label="Enter your code prompt", placeholder="def print_hello_world():")
 
         generate_btn.click(
             fn=generate_text,
-            inputs=[custom_prompt, instruction, temperature, generate_len, top_p, top_k],
+            inputs=[prompt, temperature, max_length],
             outputs=output
         )
 
-    demo.launch()
-
-if __name__ == "__main__":
-    gradio_app()
+    demo.launch()