Duplicate from replit/replit-code-v1-3b-demo

Co-authored-by: Michele Catasta <[email protected]>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Replit Code V1 3B Demo
+emoji: 🧑‍💻
+colorFrom: gray
+colorTo: orange
+sdk: gradio
+sdk_version: 3.28.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: replit/replit-code-v1-3b-demo
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,148 @@
+"""Inspired by the SantaCoder demo Huggingface space. 
+Link: https://huggingface.co/spaces/bigcode/santacoder-demo/tree/main/app.py
+"""
+
+import os
+import gradio as gr
+import torch
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
+
+REPO = "replit/replit-code-v1-3b"
+
+description = """# <h1 style="text-align: center; color: white;"><span style='color: #F26207;'> Code Completion with replit-code-v1-3b </h1>
+<span style="color: white; text-align: center;"> replit-code-v1-3b model is a 2.7B LLM trained on 20 languages from the Stack Dedup v1.2 dataset. You can click the button several times to keep completing your code.</span>"""
+
+
+token = os.environ["HUB_TOKEN"]
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+PAD_TOKEN = "<|pad|>"
+EOS_TOKEN = "<|endoftext|>"
+UNK_TOKEN = "<|unk|>"
+MAX_INPUT_TOKENS = 1024 # max tokens from context
+
+
+tokenizer = AutoTokenizer.from_pretrained(REPO, use_auth_token=token, trust_remote_code=True)
+tokenizer.truncation_side = "left" # ensures if truncate, then keep the last N tokens of the prompt going L -> R
+
+if device == "cuda":
+    model = AutoModelForCausalLM.from_pretrained(REPO, use_auth_token=token, trust_remote_code=True, low_cpu_mem_usage=True).to(device, dtype=torch.bfloat16)
+else:
+    model = AutoModelForCausalLM.from_pretrained(REPO, use_auth_token=token, trust_remote_code=True, low_cpu_mem_usage=True)
+
+model.eval()
+
+
+custom_css = """
+.gradio-container {
+    background-color: #0D1525; 
+    color:white
+}
+#orange-button {
+    background: #F26207 !important;
+    color: white;
+}
+.cm-gutters{
+    border: none !important;
+}
+"""
+
+def post_processing(prompt, completion):
+    return prompt + completion
+    # completion = "<span style='color: #499cd5;'>" + completion + "</span>"
+    # prompt = "<span style='color: black;'>" + prompt + "</span>"
+    # code_html = f"<hr><br><pre style='font-size: 14px'><code>{prompt}{completion}</code></pre><br><hr>"
+    # return code_html
+
+    
+def code_generation(prompt, max_new_tokens, temperature=0.2, seed=42, top_p=0.9, top_k=None, use_cache=True, repetition_penalty=1.0):
+
+    # truncates the prompt to MAX_INPUT_TOKENS if its too long
+    x = tokenizer.encode(prompt, return_tensors="pt", max_length=MAX_INPUT_TOKENS, truncation=True).to(device)
+    print("Prompt shape: ", x.shape) # just adding to see in the space logs in prod
+    set_seed(seed)
+    y = model.generate(x, 
+                       max_new_tokens=max_new_tokens, 
+                       temperature=temperature, 
+                       pad_token_id=tokenizer.pad_token_id, 
+                       eos_token_id=tokenizer.eos_token_id, 
+                       top_p=top_p,
+                       top_k=top_k,
+                       use_cache=use_cache,
+                       repetition_penalty=repetition_penalty
+                    )
+    completion = tokenizer.decode(y[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    completion = completion[len(prompt):]
+    return post_processing(prompt, completion)
+
+
+demo = gr.Blocks(
+    css=custom_css
+)
+
+with demo:
+    gr.Markdown(value=description)
+    with gr.Row():
+        input_col , settings_col  = gr.Column(scale=6), gr.Column(scale=6), 
+        with input_col:
+            code = gr.Code(lines=28,label='Input', value="def sieve_eratosthenes(n):")
+        with settings_col:
+            with gr.Accordion("Generation Settings", open=True):
+                max_new_tokens= gr.Slider(
+                    minimum=8,
+                    maximum=128,
+                    step=1,
+                    value=48,
+                    label="Max Tokens",
+                )
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=2.5,
+                    step=0.1,
+                    value=0.2,
+                    label="Temperature",
+                )
+                repetition_penalty = gr.Slider(
+                    minimum=1.0,
+                    maximum=1.9,
+                    step=0.1,
+                    value=1.0,
+                    label="Repetition Penalty. 1.0 means no penalty.",
+                )
+                seed = gr.Slider(
+                    minimum=0,
+                    maximum=1000,
+                    step=1,
+                    label="Random Seed"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.9,
+                    label="Top P",
+                )
+                top_k = gr.Slider(
+                    minimum=1,
+                    maximum=64,
+                    step=1,
+                    value=4,
+                    label="Top K",
+                )
+                use_cache = gr.Checkbox(
+                    label="Use Cache",
+                    value=True
+                )
+    
+    with gr.Row():
+        run = gr.Button(elem_id="orange-button", value="Generate More Code")
+
+    # with gr.Row():
+    #     # _, middle_col_row_2, _ = gr.Column(scale=1), gr.Column(scale=6), gr.Column(scale=1)
+    #     # with middle_col_row_2:
+    #     output = gr.HTML(label="Generated Code")
+
+    event = run.click(code_generation, [code, max_new_tokens, temperature, seed, top_p, top_k, use_cache, repetition_penalty], code, api_name="predict")
+
+demo.queue(max_size=40).launch()

requirements.txt

@@ -0,0 +1,6 @@
+einops
+sentencepiece
+torch
+transformers
+accelerate
+https://gradio-builds.s3.amazonaws.com/83cdcf194ba26f132eb7047c19cbadf47bc48a1c/gradio-3.28.1-py3-none-any.whl